| """Train a tiny word-level transition table. Not a neural network.""" |
| import argparse |
| import json |
| import re |
| from collections import Counter, defaultdict |
| from pathlib import Path |
|
|
|
|
| def tokenize(text): |
| return re.findall(r"[A-Za-z0-9']+|[.,!?;:]", text) |
|
|
|
|
| def train(text): |
| counts = defaultdict(Counter) |
| for line in text.splitlines(): |
| words = tokenize(line.strip()) |
| if not words: |
| continue |
| prior = '<START>' |
| for word in words: |
| counts[prior][word] += 1 |
| prior = word |
| counts[prior]['<END>'] += 1 |
| return {word: dict(next_words) for word, next_words in sorted(counts.items())} |
|
|
|
|
| if __name__ == '__main__': |
| p = argparse.ArgumentParser() |
| p.add_argument('--input', default='corpus.txt') |
| p.add_argument('--output', default='transitions.json') |
| args = p.parse_args() |
| model = train(Path(args.input).read_text(encoding='utf-8')) |
| Path(args.output).write_text(json.dumps(model, indent=2, ensure_ascii=False) + '\n', encoding='utf-8') |
| print(f'Wrote {len(model)} transition states to {args.output}') |
|
|
