util.py

import torch
import torchvision
import torchvision.transforms as transforms
from datasets import load_dataset
from torch.utils.data import Dataset
import numpy as np


class CIFAR10DataLoader:
    def __init__(self, data_dir="./data", use_augmentation=True):
        self.data_dir = data_dir

        if use_augmentation:
            self.train_transform = transforms.Compose(
                [
                    transforms.RandomHorizontalFlip(),
                    transforms.RandomCrop(32, padding=4),
                    transforms.ToTensor(),
                    transforms.Normalize(
                        (0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)
                    ),
                ]
            )
        else:
            self.train_transform = transforms.Compose(
                [
                    transforms.ToTensor(),
                    transforms.Normalize(
                        (0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)
                    ),
                ]
            )

        self.test_transform = transforms.Compose(
            [
                transforms.ToTensor(),
                transforms.Normalize(
                    (0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)
                ),
            ]
        )

        self.classes = (
            "plane",
            "car",
            "bird",
            "cat",
            "deer",
            "dog",
            "frog",
            "horse",
            "ship",
            "truck",
        )

    def load_dataset(self, train=True, download=True):
        if train:
            dataset = torchvision.datasets.CIFAR10(
                root=self.data_dir,
                train=True,
                download=download,
                transform=self.train_transform,
            )
        else:
            dataset = torchvision.datasets.CIFAR10(
                root=self.data_dir,
                train=False,
                download=download,
                transform=self.test_transform,
            )

        return dataset

    def get_dataloader(self, dataset, batch_size=1024, shuffle=True):
        """创建DataLoader"""
        return torch.utils.data.DataLoader(
            dataset, batch_size=batch_size, shuffle=shuffle, num_workers=2
        )


def get_token_id(tok, tokenizer):
    tid = tokenizer.token_to_id(tok)
    if tid is None:
        enc = tokenizer.encode(tok)
        if len(enc.ids) > 0:
            tid = enc.ids[0]
    return tid


def build_token_stream(split_name, out_path, tokenizer, EOS_ID):
    print(f"[INFO] Loading TinyStories split='{split_name}' ...")
    ds = load_dataset("roneneldan/TinyStories", split=split_name)

    ids_list = []
    print(f"[INFO] Tokenizing {split_name}...")
    from tqdm import tqdm

    for item in tqdm(ds):
        text = item.get("text", "").strip()
        if len(text) == 0:
            continue
        enc = tokenizer.encode(text, add_special_tokens=False)
        chunk = enc
        try:
            c = chunk[0]
        except TypeError:
            chunk = chunk.ids
        if EOS_ID is not None:
            if len(chunk) == 0 or chunk[-1] != EOS_ID:
                chunk += [EOS_ID for i in range(32)]
        ids_list.extend(chunk)

    arr = np.array(ids_list, dtype=np.int32)
    np.save(out_path, arr)
    print(f"[INFO] Saved {split_name} token stream to {out_path} (tokens={len(arr):,})")
    return arr


class LMDataset(Dataset):
    def __init__(self, token_array, seq_len):
        """
        token_array: np array of shape (N,)
        """
        self.tokens = torch.from_numpy(token_array).long()
        self.seq_len = seq_len

    def __len__(self):
        return len(self.tokens) - self.seq_len - 1

    def __getitem__(self, idx):
        x = self.tokens[idx : idx + self.seq_len]
        y = self.tokens[idx + 1 : idx + 1 + self.seq_len]
        return x, y


def top_k_filter(logits, k):
    if k <= 0:
        return logits
    topk_vals, topk_idx = torch.topk(logits, k)
    min_topk = topk_vals[-1]
    filtered = torch.where(
        logits < min_topk, torch.full_like(logits, float("-inf")), logits
    )
    return filtered


def top_p_filter(logits, p):
    if p >= 1.0:
        return logits
    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
    cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)

    sorted_indices_to_remove = cumulative_probs > p
    sorted_indices_to_remove[0] = False
    indices_to_remove = sorted_indices[sorted_indices_to_remove]
    filtered = logits.clone()
    filtered[indices_to_remove] = float("-inf")
    return filtered


def generate_text(
    model,
    tokenizer,
    device,
    SEQ_LEN,
    EOS_ID,
    start_string="Once",
    length=300,
    temperature=0.7,
    top_k=0,
    top_p=1.0,
    min_length=1,
):
    model.eval()
    with torch.no_grad():
        input_ids = tokenizer.encode(start_string, add_special_tokens=False)
        try:
            input_indices = list(input_ids)
        except TypeError:
            input_indices = list(input_ids.ids)
        generated_indices = list(input_indices)

        for step in range(length):
            current_indices = input_indices[-SEQ_LEN:]
            valid_len = len(current_indices)

            input_tensor = torch.full(
                (1, SEQ_LEN), fill_value=EOS_ID, dtype=torch.long
            ).to(device)
            if valid_len > 0:
                input_tensor[0, -valid_len:] = torch.tensor(
                    current_indices, dtype=torch.long
                ).to(device)

            output = model(input_tensor)
            logits = output[0, -1, :].clone()

            logits = logits / max(1e-8, temperature)

            if top_k > 0:
                logits = top_k_filter(logits, top_k)
            else:
                logits = top_k_filter(logits, np.random.randint(8, 10))
            if top_p < 1.0:
                logits = top_p_filter(logits, top_p)

            probs = torch.softmax(logits, dim=0)
            word_idx = torch.multinomial(probs, 1).item()

            if EOS_ID is not None and word_idx == EOS_ID and step >= min_length:
                generated_indices.append(word_idx)
                break
            input_indices.append(word_idx)
            generated_indices.append(word_idx)

    cut = -1
    if EOS_ID is not None and EOS_ID in generated_indices:
        cut = generated_indices.index(EOS_ID)
    if cut >= 0:
        return tokenizer.decode(generated_indices[:cut])
    else:
        return tokenizer.decode(generated_indices)