File size: 1,313 Bytes
e8c410c | 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 | """Sample from a toy word-level Markov model. No pretrained LLM is included."""
import argparse
import json
import random
import re
from pathlib import Path
def sample(model, prompt, max_new_words, seed):
rng = random.Random(seed)
tokens = re.findall(r"[A-Za-z0-9']+|[.,!?;:]", prompt)
state = tokens[-1] if tokens else '<START>'
output = list(tokens)
for _ in range(max_new_words):
options = model.get(state) or model.get('<START>', {})
if not options:
break
candidates, weights = zip(*options.items())
word = rng.choices(candidates, weights=weights, k=1)[0]
if word == '<END>':
break
output.append(word)
state = word
text = ' '.join(output)
return re.sub(r'\s+([.,!?;:])', r'\1', text)
if __name__ == '__main__':
p = argparse.ArgumentParser()
p.add_argument('--model', default='transitions.json')
p.add_argument('--prompt', default='The future of AI')
p.add_argument('--max-new-words', type=int, default=35)
p.add_argument('--seed', type=int, default=7)
a = p.parse_args()
if a.max_new_words < 0:
p.error('--max-new-words must be >= 0')
model = json.loads(Path(a.model).read_text(encoding='utf-8'))
print(sample(model, a.prompt, a.max_new_words, a.seed))
|