| """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.") |
| |
| 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) |
| |
| 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 |
|
|
