| """Sample a continuation from DWSLM-256M (a base model: it continues text, it doesn't chat). |
| |
| python generate.py "The three primary colors are" |
| python generate.py "def fibonacci(n):" --temperature 0 --max-new 80 |
| |
| There's no KV cache: each new token re-runs the whole prefix, which is fine at these lengths on a GPU. |
| """ |
|
|
| import argparse |
|
|
| import torch |
| from dwslm import EOS_ID, context, load, next_logits |
|
|
| UNFINISHED = "\ufffd" |
|
|
|
|
| def main() -> None: |
| p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| p.add_argument("prompt") |
| p.add_argument("--max-new", type=int, default=64) |
| p.add_argument("--temperature", type=float, default=0.7, help="0 = greedy") |
| p.add_argument("--top-p", type=float, default=0.9) |
| p.add_argument("--repetition-penalty", type=float, default=1.1, help="1 = off") |
| p.add_argument("--repo", default=".", help="folder holding config.json and model.safetensors") |
| a = p.parse_args() |
|
|
| model, tok = load(a.repo) |
| ids = [EOS_ID, *tok.encode(a.prompt).ids] |
| start, shown = len(ids), "" |
| print(a.prompt, end="", flush=True) |
| for _ in range(min(a.max_new, context(model) - len(ids))): |
| z = next_logits(model, [ids])[0] |
| nxt = _sample(z, ids, a) |
| if nxt == EOS_ID: |
| break |
| ids.append(nxt) |
|
|
| |
| text = tok.decode(ids[start:]).rstrip(UNFINISHED) |
| print(text[len(shown) :], end="", flush=True) |
| shown = text |
| print(tok.decode(ids[start:])[len(shown) :]) |
|
|
|
|
| def _sample(z: torch.Tensor, seen: list[int], a: argparse.Namespace) -> int: |
| if a.repetition_penalty != 1.0: |
| idx = torch.tensor(sorted(set(seen)), device=z.device) |
| prev = z[idx] |
| z[idx] = torch.where(prev > 0, prev / a.repetition_penalty, prev * a.repetition_penalty) |
| if a.temperature <= 0: |
| return int(z.argmax()) |
|
|
| probs = (z / a.temperature).softmax(-1) |
| p, order = probs.sort(descending=True) |
| p = p * (p.cumsum(-1) - p < a.top_p) |
| return int(order[torch.multinomial(p / p.sum(), 1)]) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
