File size: 1,956 Bytes
c21e887 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | import re
import os
import json
from abc import ABC, abstractmethod
from typing import Iterable
import tiktoken
_PATTERN = re.compile(r'([,.:;?_!"()\']|--|\s)')
def build_vocab(text: str) -> dict[str, int]:
"""Build a vocab dict {word -> int} from raw text."""
tokens = [t for t in _PATTERN.split(text.lower()) if t and not t.isspace()]
all_words: list[str] = sorted(set(tokens))
all_words.extend(["<unk>", "<eos>"])
return {word: idx for idx, word in enumerate(all_words)}
def save_vocab(vocab: dict[str, int], file_path: str) -> None:
os.makedirs(file_path, exist_ok=True)
with open(os.path.join(file_path, "vocab.json"), "w") as f:
json.dump(vocab, f, indent=2)
class BaseTokenizer(ABC):
@abstractmethod
def encode(self, text: str) -> list[int]: ...
@abstractmethod
def decode(self, tokens: Iterable[int]) -> str: ...
class SimpleTokenizer(BaseTokenizer):
"""Vanilla tokenizer that builds a str→int map over the whole dataset."""
def __init__(self, vocab: dict[str, int]) -> None:
self.str_to_int = vocab
self.int_to_str = {v: k for k, v in vocab.items()}
def encode(self, text: str) -> list[int]:
tokens = [t for t in _PATTERN.split(text.lower()) if t and not t.isspace()]
ids = [self.str_to_int.get(t, self.str_to_int["<unk>"]) for t in tokens]
ids.append(self.str_to_int["<eos>"])
return ids
def decode(self, tokens: Iterable[int]) -> str:
return " ".join(self.int_to_str.get(t, "<unk>") for t in tokens)
class TikTokenizer(BaseTokenizer):
"""BPE tokenizer used in GPT-2, via the tiktoken library."""
def __init__(self, model_type: str) -> None:
self.tokenizer = tiktoken.get_encoding(model_type)
def encode(self, text: str) -> list[int]:
return self.tokenizer.encode(text)
def decode(self, tokens: Iterable[int]) -> str:
return self.tokenizer.decode(list(tokens))
|