| """Talk to brAIAIn from the command line. Standard library only. |
| |
| python generate.py "What is the meaning of life?" |
| |
| Uses brAIAIn-latest.json (the newest version) if present, otherwise the release's own |
| weights file in this folder (for example brAIAIn-2.67-turbo.json). |
| """ |
| import json, random, re, sys |
|
|
| import os |
| _here = os.path.dirname(os.path.abspath(__file__)) |
| _jsons = sorted(f for f in os.listdir(_here) if f.startswith("brAIAIn") and f.endswith(".json")) |
| _file = "brAIAIn-latest.json" if "brAIAIn-latest.json" in _jsons else (_jsons[0] if _jsons else "brAIAIn-0.json") |
| m = json.load(open(os.path.join(_here, _file), encoding="utf-8")) |
| BOS, EOS = "<s>", "</s>" |
|
|
|
|
| def pick(table): |
| keys = sorted(table) |
| r = random.random() * sum(table.values()) |
| for k in keys: |
| r -= table[k] |
| if r <= 0: |
| return k |
| return keys[-1] |
|
|
|
|
| def answer(prompt="", max_words=40): |
| known = {re.sub(r"[^a-z0-9']", "", w.lower()): w for w in m["bi"] if w != BOS} |
| words = [re.sub(r"[^a-z0-9']", "", w.lower()) for w in prompt.split()] |
| starts = [known[w] for w in words if w in known and len(w) > 3] |
| a, b, out = BOS, BOS, [] |
| if starts: |
| b = random.choice(starts) |
| out.append(b) |
| while len(out) < max_words: |
| k = a + " " + b |
| nxt = m["tri"].get(k) or m["bi"].get(b) |
| if not nxt: |
| break |
| c = pick(nxt) |
| if c == EOS: |
| break |
| out.append(c) |
| a, b = b, c |
| s = " ".join(out) |
| return s[:1].upper() + s[1:] |
|
|
|
|
| print(answer(" ".join(sys.argv[1:]))) |
|
|
