File size: 4,390 Bytes
20a371d | 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 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | """Load a standalone Sprout bundle and stream instruction-tuned stories."""
import hashlib
import json
import math
from pathlib import Path
import torch
import tokenizer
from model import Sprout
TEMPLATE = "<|user|>\n{instruction}<|end|>\n<|assistant|>\n"
FIELDS = ("Summary", "Words", "Random sentence", "Features")
def instruction(fields):
return "Write a short story.\n" + "\n".join(
f"{key}: {fields[key]}" for key in FIELDS if fields.get(key))
def prompt(tok, fields):
return tok.encode(TEMPLATE.format(instruction=instruction(fields)))
def digest(path):
with open(path, "rb") as f:
return hashlib.file_digest(f, "sha256").hexdigest()
def load(source, device="cpu", revision=None):
"""source is an exported directory or Hub model ID; revision may pin a commit."""
from safetensors.torch import load_model
path = Path(source)
if not path.is_dir():
from huggingface_hub import snapshot_download
path = Path(snapshot_download(str(source), revision=revision,
allow_patterns=["config.json", "model.safetensors", "tokenizer.model"]))
cfg = json.loads((path / "config.json").read_text())
if cfg["format_version"] != 1 or cfg["template"] != TEMPLATE or cfg["fields"] != list(FIELDS):
raise ValueError("Unsupported model format or instruction template.")
for name, expected in cfg["sha256"].items():
if name not in ("model.safetensors", "tokenizer.model") or digest(path / name) != expected:
raise ValueError(f"Bundle checksum mismatch: {name}")
if set(cfg["sha256"]) != {"model.safetensors", "tokenizer.model"}:
raise ValueError("Bundle must checksum both weights and tokenizer.")
tok = tokenizer.load(path / "tokenizer.model")
if len(tok.vocab) != cfg["model"]["vocab"] or tok.special != cfg["special_tokens"]:
raise ValueError("Tokenizer does not match the model configuration.")
model = Sprout(**cfg["model"])
load_model(model, path / "model.safetensors", strict=True)
return model.to(device).eval(), tok, cfg
def stream(model, tok, fields, temperature=0.8, top_k=50, max_tokens=256):
"""Yield plain text; UTF-8 bytes are buffered until each character is complete."""
import codecs
if set(fields) - set(FIELDS) or any(not isinstance(v, str) for v in fields.values()):
raise ValueError("Use the supported story fields with text values.")
if sum(len(v) for v in fields.values()) > 4096:
raise ValueError("Story instructions are too long. Shorten them and try again.")
if any(s in v for v in fields.values() for s in tokenizer.SPECIALS):
raise ValueError("Remove reserved tokens such as <|user|> from the story fields.")
# Each field occupies one metadata line in the training template.
fields = {key: " ".join(value.split()) for key, value in fields.items()}
if not any(fields.values()):
raise ValueError("Enter a story idea or an optional story constraint.")
if not math.isfinite(float(temperature)) or not 0 <= temperature <= 2:
raise ValueError("Temperature must be between 0 and 2; 0 selects greedy decoding.")
if not isinstance(top_k, int) or not 1 <= top_k <= len(tok.vocab):
raise ValueError("Top-k must be a positive integer within the vocabulary.")
if not isinstance(max_tokens, int) or not 1 <= max_tokens <= model.ctx:
raise ValueError(f"Maximum tokens must be between 1 and {model.ctx}.")
prefix = prompt(tok, fields)
# Leave room for a story, and keep the entire instruction in view.
if len(prefix) > model.ctx - 32:
raise ValueError("Story instructions use too much context. Shorten them and try again.")
limit = min(max_tokens, model.ctx - len(prefix) + 1)
x = torch.tensor([prefix], device=next(model.parameters()).device)
decoder = codecs.getincrementaldecoder("utf-8")(errors="replace")
text = ""
yield text
tokens = model.generate_tokens(x, limit, temperature=temperature, top_k=top_k,
stop=tok.special["<|end|>"])
for token in tokens:
idx = token.item()
if idx in tok.special.values():
break
part = decoder.decode(tok.vocab[idx])
if part:
text += part
yield text
tail = decoder.decode(b"", final=True)
if tail:
yield text + tail
|