| """Run text generation with the from-scratch TinyStories Transformer. |
| |
| Setup (uv): |
| uv venv --python 3.12 # torch~=2.6.0 has no wheels yet for newer Pythons |
| uv pip install "cs336_basics @ git+https://github.com/zhL-d/llm-from-scratch.git@2b1593c1bc049dde93d332b9f5a6cf9d57575744" |
| uv pip install huggingface_hub |
| uv run generate.py "Once upon a time" |
| |
| This script only depends on the `cs336_basics` package (the from-scratch |
| implementation, pinned to the exact commit this checkpoint was produced |
| under) and `huggingface_hub` (to pull the tokenizer files from the |
| dataset repo). The model weights (`model.pt`) and `config.json` in this |
| repo are loaded directly from disk. |
| """ |
| import json |
| import sys |
| from pathlib import Path |
|
|
| import torch |
| from huggingface_hub import hf_hub_download |
|
|
| from cs336_basics.transformer_lm import TransformerLM |
| from cs336_basics.tokenizer import Tokenizer |
| from cs336_basics.decoding import Decoding |
|
|
| HERE = Path(__file__).parent |
|
|
|
|
| def load_tokenizer(config: dict) -> Tokenizer: |
| vocab_path = hf_hub_download( |
| repo_id=config["tokenizer_repo"], |
| filename=config["vocab_file"], |
| repo_type="dataset", |
| ) |
| merges_path = hf_hub_download( |
| repo_id=config["tokenizer_repo"], |
| filename=config["merges_file"], |
| repo_type="dataset", |
| ) |
| return Tokenizer.from_files( |
| vocab_filepath=vocab_path, |
| merges_filepath=merges_path, |
| special_tokens=config["special_tokens"], |
| ) |
|
|
|
|
| def load_model(config: dict) -> TransformerLM: |
| model = TransformerLM( |
| vocab_size=config["vocab_size"], |
| context_length=config["context_length"], |
| num_layers=config["num_layers"], |
| d_model=config["d_model"], |
| num_heads=config["num_heads"], |
| d_ff=config["d_ff"], |
| rope_theta=config["rope_theta"], |
| ) |
| state = torch.load(HERE / "model.pt", map_location="cpu") |
| model.load_state_dict(state["model"]) |
| model.eval() |
| return model |
|
|
|
|
| def generate( |
| prompt: str, |
| max_tokens: int = 256, |
| temperature: float = 0.7, |
| topp_threshold: float = 0.9, |
| ) -> str: |
| config = json.loads((HERE / "config.json").read_text()) |
|
|
| tokenizer = load_tokenizer(config) |
| model = load_model(config) |
|
|
| prompt_ids = tokenizer.encode(prompt) |
| eos_id = tokenizer.encode(config["special_tokens"][0])[0] |
|
|
| generated_ids = Decoding( |
| model=model, |
| prompt_ids=prompt_ids, |
| max_tokens=max_tokens, |
| temperature=temperature, |
| topp_threshold=topp_threshold, |
| eos_id=eos_id, |
| context_length=config["context_length"], |
| ) |
|
|
| return tokenizer.decode(generated_ids[0].tolist()) |
|
|
|
|
| if __name__ == "__main__": |
| prompt = sys.argv[1] if len(sys.argv) > 1 else "Once upon a time" |
| print(generate(prompt)) |
|
|
