| """Continue a text prompt with MP-01 100M base. |
| |
| Usage (from this folder): |
| python run.py "The history of the printing press" |
| python run.py "Once upon a time" --tokens 200 --temperature 0.7 |
| |
| Runs on a CUDA GPU when one is available, otherwise on the CPU. |
| """ |
| import argparse |
| import json |
| import sys |
| from pathlib import Path |
|
|
| import torch |
| from safetensors.torch import load_file |
| from tokenizers import Tokenizer |
|
|
| HERE = Path(__file__).resolve().parent |
| sys.path.insert(0, str(HERE)) |
| from vbq.model import VBQ, VBQConfig |
|
|
|
|
| def main(): |
| p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| p.add_argument("prompt", nargs="?", default="The history of the printing press") |
| p.add_argument("--tokens", type=int, default=150, help="new tokens to generate") |
| p.add_argument("--temperature", type=float, default=0.8) |
| p.add_argument("--top-k", type=int, default=50) |
| p.add_argument("--top-p", type=float, default=0.95) |
| p.add_argument("--repetition-penalty", type=float, default=1.1) |
| p.add_argument("--seed", type=int, default=None) |
| a = p.parse_args() |
| if a.seed is not None: |
| torch.manual_seed(a.seed) |
|
|
| device = "cuda" if torch.cuda.is_available() else "cpu" |
| cfg = VBQConfig(**json.loads((HERE / "config.json").read_text())["model_cfg"]) |
| model = VBQ(cfg) |
| model.load_state_dict(load_file(str(HERE / "model.safetensors"))) |
| model.to(device=device, dtype=torch.bfloat16 if device == "cuda" else torch.float32).eval() |
|
|
| tok = Tokenizer.from_file(str(HERE / "tokenizer.json")) |
| ids = tok.encode(a.prompt, add_special_tokens=False).ids |
| idx = torch.tensor([ids], device=device) |
| print(a.prompt, end="", flush=True) |
| new, shown = [], "" |
| with torch.no_grad(): |
| for t in model.generate_cached(idx, max_new_tokens=a.tokens, temperature=a.temperature, top_k=a.top_k, |
| top_p=a.top_p, eos_id=0, repetition_penalty=a.repetition_penalty): |
| new.append(t) |
| text = tok.decode(new) |
| print(text[len(shown):], end="", flush=True) |
| shown = text |
| print() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
