File size: 1,098 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 | """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}')
|