| |
| """test.py — Micro Distillery model chat. |
| |
| Note: the model has no explicit EOS token. Generation is bounded by max_tokens. |
| BOS (space) is a soft end-of-trajectory signal inside the browser trainer only. |
| """ |
| import json, os, random, sys |
|
|
| MODEL_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "model.json") |
|
|
| def load_model(path=MODEL_PATH): |
| with open(path, "r", encoding="utf-8") as f: |
| return json.load(f) |
|
|
| def sample_next(transitions, token, rng, temperature=1.0, top_k=0): |
| bucket = transitions.get(token) |
| if not bucket: return None |
| items = [(k, v) for k, v in bucket.items() if v > 0] |
| if not items: return None |
| if top_k and len(items) > top_k: |
| items = sorted(items, key=lambda x: -x[1])[:top_k] |
| tokens = [t for t, _ in items] |
| probs = [p for _, p in items] |
| if temperature <= 0: |
| return max(items, key=lambda x: x[1])[0] |
| inv_t = 1.0 / temperature |
| weights = [p ** inv_t for p in probs] |
| total = sum(weights) |
| if total <= 0: return None |
| r = rng.random() * total |
| acc = 0.0 |
| for t, w in zip(tokens, weights): |
| acc += w |
| if r <= acc: return t |
| return tokens[-1] |
|
|
| def generate(model, prompt, max_tokens=32, temperature=0.9, top_k=0): |
| transitions = model.get("transitions", {}) |
| seed = int(model.get("harmonic_seed", 1337)) & 0xFFFFFFFF |
| rng = random.Random(seed) |
| prefix = prompt or "" |
| current = prefix[-1] if prefix else " " |
| generated = "" |
| for _ in range(max_tokens): |
| nxt = sample_next(transitions, current, rng, temperature, top_k) |
| if nxt is None: break |
| generated += nxt |
| current = nxt |
| return prefix + generated |
|
|
| def main(temperature=0.9): |
| print("=== Micro Distillery Chat (toy bigram model) ===") |
| if not os.path.exists(MODEL_PATH): |
| print("Error: " + MODEL_PATH + " not found.", file=sys.stderr) |
| sys.exit(1) |
| model = load_model() |
| print("Format: " + str(model.get("format"))) |
| print("Vocab: " + str(len(model.get("vocab", [])))) |
| print("States: " + str(len(model.get("transitions", {})))) |
| print("Temperature: " + str(temperature)) |
| print("Type /quit to exit.\n") |
| while True: |
| try: prompt = input("You: ").strip() |
| except (EOFError, KeyboardInterrupt): break |
| if not prompt or prompt.lower() in ("/quit", "/exit"): break |
| print("Bot: " + generate(model, prompt, temperature=temperature)) |
|
|
| if __name__ == "__main__": |
| temperature = 0.9 |
| prompts = [] |
| i = 1 |
| while i < len(sys.argv): |
| if sys.argv[i] == "--temperature" and i + 1 < len(sys.argv): |
| temperature = float(sys.argv[i + 1]); i += 2 |
| else: |
| prompts.append(sys.argv[i]); i += 1 |
| if prompts: |
| model = load_model() |
| print(generate(model, " ".join(prompts), temperature=temperature)) |
| else: |
| main(temperature) |
|
|
