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