| |
| """ |
| test.py - Micro Distillery model chat wrapper. |
| |
| Loads model.json (a GRPO+VAE-distilled bigram model) and provides |
| a simple REPL chat interface. |
| |
| Usage: |
| python3 test.py # interactive chat |
| python3 test.py hello # one-shot from prompt |
| python3 test.py --temperature 1.2 # one-shot with custom temperature |
| python3 test.py # (interactive also supports --temperature) |
| """ |
| import json |
| import os |
| import random |
| import 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): |
| bucket = transitions.get(token) |
| if not bucket: |
| return None |
| tokens = list(bucket.keys()) |
| inv_t = 1.0 / max(temperature, 1e-3) |
| weights = [max(bucket[t], 1e-12) ** inv_t for t in tokens] |
| 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): |
| transitions = model.get("transitions", {}) |
| bos = model.get("bos_token", "<|endoftext|>") |
| eos = model.get("eos_token", "<|endoftext|>") |
| seed = int(model.get("harmonic_seed", 1337)) & 0xFFFFFFFF |
| rng = random.Random(seed) |
|
|
| if prompt and prompt.strip(): |
| current = prompt.strip()[-1] |
| prefix = prompt.strip() |
| else: |
| current = bos |
| prefix = "" |
|
|
| generated = "" |
| for _ in range(max_tokens): |
| nxt = sample_next(transitions, current, rng, temperature=temperature) |
| if nxt is None or nxt == eos: |
| break |
| generated += nxt |
| current = nxt |
| return prefix + generated |
|
|
|
|
| def main(temperature=0.9): |
| print("=== Micro Distillery Model Chat ===") |
| 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 size: " + str(len(model.get("vocab", [])))) |
| print("Transitions: " + str(len(model.get("transitions", {})))) |
| print("Seed: " + str(model.get("harmonic_seed"))) |
| print("Temperature: " + str(temperature)) |
| print("Type /quit to exit.\n") |
| while True: |
| try: |
| prompt = input("You: ").strip() |
| except (EOFError, KeyboardInterrupt): |
| print("\nBye!") |
| break |
| if not prompt: |
| continue |
| if prompt.lower() in ("/quit", "/exit", "quit", "exit"): |
| break |
| out = generate(model, prompt, temperature=temperature) |
| print("Bot: " + out) |
|
|
|
|
| if __name__ == "__main__": |
| args = sys.argv[1:] |
| temperature = 0.9 |
| prompts = [] |
| i = 0 |
| while i < len(args): |
| if args[i] == "--temperature": |
| if i + 1 < len(args): |
| try: |
| temperature = float(args[i + 1]) |
| except ValueError: |
| print("Warning: invalid --temperature value, using 0.9", file=sys.stderr) |
| i += 2 |
| else: |
| print("Warning: --temperature requires a value, using 0.9", file=sys.stderr) |
| i += 1 |
| elif args[i] in ("--help", "-h"): |
| print("Usage: python3 test.py [prompt...] [--temperature T]") |
| sys.exit(0) |
| elif args[i].startswith("--"): |
| i += 1 |
| else: |
| prompts.append(args[i]) |
| i += 1 |
| if prompts: |
| model = load_model() |
| print(generate(model, " ".join(prompts), temperature=temperature)) |
| else: |
| main(temperature=temperature) |
|
|
