Spaces:
Sleeping
Sleeping
Premchandyadav369
feat: deploy interactive Q-TensorFormer demo with trained weights and Gradio 4 UI
67ded87 | """ | |
| Data loading and preprocessing. | |
| Supported datasets: | |
| - WikiText-2 (char-level and word-level) | |
| - WikiText-103 | |
| - Custom text files | |
| - Synthetic random data (debugging) | |
| Tokenization: character-level by default. Simple, deterministic, no external deps. | |
| """ | |
| import torch | |
| from torch.utils.data import Dataset, DataLoader | |
| from typing import Optional, Tuple, Dict, Union | |
| from collections import Counter | |
| from .tokenizer import QTensorFormerTokenizer | |
| class CharTokenizer: | |
| """Character-level tokenizer. Vocabulary built from data.""" | |
| def __init__(self, min_freq: int = 1): | |
| self.min_freq = min_freq | |
| self.char_to_idx: Dict[str, int] = {} | |
| self.idx_to_char: Dict[int, str] = {} | |
| self.vocab_size = 0 | |
| self.special_tokens = { | |
| "<pad>": 0, | |
| "<bos>": 1, | |
| "<eos>": 2, | |
| "<unk>": 3, | |
| } | |
| def fit(self, texts: list[str]): | |
| """Build vocabulary from texts.""" | |
| char_counts = Counter() | |
| for text in texts: | |
| char_counts.update(text) | |
| # Special tokens first | |
| self.char_to_idx = dict(self.special_tokens) | |
| # Freq-filtered chars | |
| idx = len(self.special_tokens) | |
| for char, count in sorted(char_counts.items(), key=lambda kv: (-kv[1], kv[0])): | |
| if count >= self.min_freq: | |
| self.char_to_idx[char] = idx | |
| idx += 1 | |
| self.idx_to_char = {v: k for k, v in self.char_to_idx.items()} | |
| self.vocab_size = len(self.char_to_idx) | |
| def encode(self, text: str, add_bos: bool = True, | |
| add_eos: bool = True, max_len: int = None) -> list[int]: | |
| """Convert text to token indices.""" | |
| tokens = [] | |
| if add_bos: | |
| tokens.append(self.special_tokens["<bos>"]) | |
| for ch in text: | |
| tokens.append(self.char_to_idx.get(ch, self.special_tokens["<unk>"])) | |
| if add_eos: | |
| tokens.append(self.special_tokens["<eos>"]) | |
| if max_len is not None: | |
| if len(tokens) > max_len: | |
| tokens = tokens[:max_len] | |
| else: | |
| tokens.extend([self.special_tokens["<pad>"]] * (max_len - len(tokens))) | |
| return tokens | |
| def decode(self, indices: list[int], skip_special: bool = True) -> str: | |
| """Convert indices back to text.""" | |
| chars = [] | |
| for idx in indices: | |
| ch = self.idx_to_char.get(int(idx), "?") | |
| if skip_special and int(idx) in self.special_tokens.values(): | |
| continue | |
| chars.append(ch) | |
| return "".join(chars) | |
| def save(self, path: str): | |
| torch.save({ | |
| "char_to_idx": self.char_to_idx, | |
| "idx_to_char": self.idx_to_char, | |
| "vocab_size": self.vocab_size, | |
| "special_tokens": self.special_tokens, | |
| }, path) | |
| def load(cls, path: str) -> "CharTokenizer": | |
| data = torch.load(path, map_location="cpu", weights_only=False) | |
| tok = cls() | |
| tok.char_to_idx = data["char_to_idx"] | |
| tok.idx_to_char = data["idx_to_char"] | |
| tok.vocab_size = data["vocab_size"] | |
| tok.special_tokens = data["special_tokens"] | |
| return tok | |
| class TextDataset(Dataset): | |
| """ | |
| Causal language modeling dataset. | |
| Splits text into overlapping sequences of length seq_len. | |
| Target = input shifted by 1 (next-token prediction). | |
| """ | |
| def __init__(self, texts: list[str], tokenizer: CharTokenizer, | |
| seq_len: int = 128, stride: int = None): | |
| self.seq_len = seq_len | |
| self.stride = stride or seq_len // 2 | |
| # Tokenize all texts | |
| all_tokens = [] | |
| for text in texts: | |
| all_tokens.extend(tokenizer.encode(text, add_bos=False, add_eos=True)) | |
| self.tokens = torch.tensor(all_tokens, dtype=torch.long) | |
| # Compute valid starting positions | |
| self.n_samples = max(0, (len(self.tokens) - seq_len - 1) // self.stride + 1) | |
| def __len__(self): | |
| return self.n_samples | |
| def __getitem__(self, idx): | |
| start = idx * self.stride | |
| end = start + self.seq_len | |
| x = self.tokens[start:end] | |
| y = self.tokens[start + 1:end + 1] | |
| assert len(x) == len(y) == self.seq_len, f"len={len(x)} at idx={idx}" | |
| return x, y | |
| def load_wikitext2(tokenizer: CharTokenizer = None, | |
| seq_len: int = 128, | |
| batch_size: int = 16) -> Tuple[DataLoader, DataLoader, DataLoader, CharTokenizer]: | |
| """ | |
| Load WikiText-2 with char-level tokenization. | |
| Returns: | |
| train_loader, val_loader, test_loader, tokenizer | |
| """ | |
| try: | |
| from datasets import load_dataset | |
| except ImportError: | |
| raise ImportError("pip install datasets") | |
| ds = load_dataset("wikitext", "wikitext-2-raw-v1") | |
| # Filter empty lines | |
| train_texts = [t for t in ds["train"]["text"] if t.strip()] | |
| val_texts = [t for t in ds["validation"]["text"] if t.strip()] | |
| test_texts = [t for t in ds["test"]["text"] if t.strip()] | |
| if tokenizer is None: | |
| tokenizer = CharTokenizer() | |
| tokenizer.fit(train_texts) | |
| train_ds = TextDataset(train_texts, tokenizer, seq_len) | |
| val_ds = TextDataset(val_texts, tokenizer, seq_len) | |
| test_ds = TextDataset(test_texts, tokenizer, seq_len) | |
| train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, | |
| num_workers=0, drop_last=True) | |
| val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False, num_workers=0) | |
| test_loader = DataLoader(test_ds, batch_size=batch_size, shuffle=False, num_workers=0) | |
| return train_loader, val_loader, test_loader, tokenizer | |
| def load_synthetic_data(vocab_size: int = 5000, seq_len: int = 128, | |
| n_samples: int = 2000, batch_size: int = 16): | |
| """Synthetic random data for debugging.""" | |
| class _SynthDataset(Dataset): | |
| def __init__(self, n, vocab, slen): | |
| self.data = torch.randint(1, vocab, (n, slen + 1)) | |
| def __len__(self): | |
| return len(self.data) | |
| def __getitem__(self, i): | |
| return self.data[i, :-1], self.data[i, 1:] | |
| ds = _SynthDataset(n_samples, vocab_size, seq_len) | |
| return DataLoader(ds, batch_size=batch_size, shuffle=True, num_workers=0) | |