File size: 3,870 Bytes
385fc44 | 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 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | #!/usr/bin/env python3
"""
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)
|