File size: 2,924 Bytes
9a452ca | 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 | #!/usr/bin/env python3
"""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)
|