File size: 2,273 Bytes
51e1e0c | 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 | """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 # noqa: E402
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) # decode the whole tail so multi-byte characters print correctly
print(text[len(shown):], end="", flush=True)
shown = text
print()
if __name__ == "__main__":
main()
|