File size: 2,214 Bytes
ca76ef5 | 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 | """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)
# print only what's new; a byte-level token can end mid-character ("\ufffd") until the rest arrives
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()
|