NotShrirang/QuillGPT

1import json2import os3from typing import Iterable4import torch5 6class Tokenizer:7    def __init__(self, data_path: str = None):8        self.config = None9        self.stoi = None10        self.itos = None11        self.vocab_size = None12        if data_path:13            self.data = self.load_data(data_path)14        else:15            self.data = None16    17    def from_pretrained(self, config_path: str):18        with open(config_path) as f:19            config = json.load(f)20        self.config = config21        if 'encode' not in config:22            raise ValueError("Config file must contain an 'encode' key.")23        if 'decode' not in config:24            raise ValueError("Config file must contain a 'decode' key.")25        if 'vocab_size' not in config:26            raise ValueError("Config file must contain a 'vocab_size' key.")27        stoi = config['encode']28        self.stoi = {k: int(v) for k, v in stoi.items()}29        itos = config['decode']30        self.itos = {int(k): v for k, v in itos.items()}31        self.vocab_size = config['vocab_size']32        return self33    34    def load_data(self, path: str) -> str:35        if not os.path.exists(path):36            raise FileNotFoundError("File not found.")37        if not path.endswith('.txt'):38            raise ValueError("File must be a text file.")39        with open(path, 'r', encoding='utf-8') as f:40            text = f.read()41        chars = sorted(list(set(text)))42        vocab_size = len(chars)43        stoi = {ch: i for i, ch in enumerate(chars)}44        itos = {i: ch for i, ch in enumerate(chars)}45        self.config = {"vocab_size": vocab_size, "encode": stoi, "decode": itos}46        self.stoi = stoi47        self.itos = itos48        data = torch.tensor(self(text), dtype=torch.long)49        n = int(0.9*len(data))50        train_data = data[:n]51        val_data = data[n:]52        self.train_data = train_data53        self.val_data = val_data54        self.vocab_size = vocab_size55        return text56 57    def __repr__(self) -> str:58        if self.config:59            return f"Tokenizer(config={self.config})"60        else:61            return f"Tokenizer()"62    63    def __str__(self) -> str:64        if self.config:65            return f"Tokenizer(config_path={self.config})"66        else:67            return f"Tokenizer()"68    69    def __len__(self) -> int:70        return len(self.stoi)71    72    def __getitem__(self, key: str) -> int:73        return self.stoi[key]74    75    def __contains__(self, key: str) -> bool:76        return key in self.stoi77    78    def __iter__(self):79        return iter(self.stoi)80    81    def __reversed__(self):82        return reversed(self.stoi)83    84    def keys(self):85        return self.stoi.keys()86    87    def values(self):88        return self.stoi.values()89    90    def items(self):91        return self.stoi.items()92    93    def __call__(self, *args, **kwds) -> list[int]:94        return self.encode(*args, **kwds)95 96    def encode(self, s: str | list[str]) -> list[int]:97        if isinstance(s, str):98            return [self.stoi[c] for c in s]99        elif isinstance(s, list):100            return [[self.stoi[i] for i in c] for c in s]101        else:102            raise ValueError("Input must be a string or a list of strings.")103 104    def decode(self, l: list[int]) -> str:105        if isinstance(l[0], int):106            return ''.join([self.itos[i] for i in l])107        elif isinstance(l[0], Iterable):108            return [''.join([self.itos[i] for i in c]) for c in l]109        else:110            raise ValueError("Input must be a list of integers or a list of list of integers.")111    112    def save_pretrained(self, path: str) -> str:113        with open(path + 'vocab.json', 'w') as f:114            json.dump(self.config, f)115        return "Tokenizer saved at {}.".format(path)116