import argparse import os from pathlib import Path import torch from train import TinyTransformerLM, build_vocab, encode_text, make_batch def load_or_create_model(model_path, text, block_size, n_embd, n_head, n_layer): requested_config = { "vocab_size": None, "block_size": block_size, "n_embd": n_embd, "n_head": n_head, "n_layer": n_layer, } if model_path.exists(): checkpoint = torch.load(model_path, map_location="cpu") saved_stoi = checkpoint["stoi"] fresh_stoi, fresh_itos = build_vocab(text) saved_config = checkpoint["config"] config_matches = all(saved_config.get(k) == v for k, v in requested_config.items() if v is not None) if set(saved_stoi) == set(fresh_stoi) and config_matches: model = TinyTransformerLM(**saved_config) model.load_state_dict(checkpoint["model"]) return model, saved_stoi, {int(k): v for k, v in checkpoint["itos"].items()}, saved_config backup = model_path.with_suffix(".old-vocab.pt") model_path.replace(backup) print(f"old vocabulary checkpoint moved to {backup}") stoi, itos = build_vocab(text) config = { "vocab_size": len(stoi), "block_size": block_size, "n_embd": n_embd, "n_head": n_head, "n_layer": n_layer, } model = TinyTransformerLM(**config) return model, stoi, itos, config def limit_memory_file(data_path, keep_tail_chars=200_000): if not data_path.exists(): return text = data_path.read_text(encoding="utf-8") if len(text) <= keep_tail_chars: return data_path.write_text(text[-keep_tail_chars:], encoding="utf-8") def normalize_text(text): return "".join(ch.lower() if ch.isalpha() else ch for ch in text) def score_snippet(query, snippet): query_chars = set(ch for ch in normalize_text(query) if not ch.isspace()) snippet_chars = set(ch for ch in normalize_text(snippet) if not ch.isspace()) if not query_chars or not snippet_chars: return 0 overlap = len(query_chars & snippet_chars) return overlap / len(query_chars | snippet_chars) def retrieve_context(user_text, memory_text, max_examples=3): blocks = [b.strip() for b in memory_text.split("\n\n") if b.strip()] scored = [] for block in blocks: if "USER:" in block and "AI:" in block: scored.append((score_snippet(user_text, block), block)) scored.sort(key=lambda item: item[0], reverse=True) chosen = [block for score, block in scored[:max_examples] if score > 0] return "\n\n".join(chosen) @torch.no_grad() def generate_once(model, prompt, stoi, itos, max_new_tokens, temperature): fallback = next(iter(stoi.values())) idx = torch.tensor([[stoi.get(ch, fallback) for ch in prompt]], dtype=torch.long) model.eval() for _ in range(max_new_tokens): idx_cond = idx[:, -model.block_size :] logits, _ = model(idx_cond) logits = logits[:, -1, :] / temperature probs = torch.softmax(logits, dim=-1) next_id = torch.multinomial(probs, num_samples=1) idx = torch.cat((idx, next_id), dim=1) if itos[int(next_id)] == "\n" and idx.shape[1] > len(prompt) + 20: break return "".join(itos[int(i)] for i in idx[0])[len(prompt) :] def score_reply(reply, user_text): stripped = reply.strip() if not stripped: return -100 score = 0.0 lowered = stripped.lower() if "user:" in lowered: score -= 6 if "ai:" in lowered: score -= 6 if stripped.count("\n") > 2: score -= 2 if len(set(stripped)) < 4: score -= 2 words = [word for word in stripped.split() if word] if words: unique_ratio = len(set(words)) / len(words) score += unique_ratio * 3 if len(stripped) < 4: score -= 2 if len(stripped) > 220: score -= 1 if user_text and user_text.lower().strip() in lowered: score -= 2 if any(ch.isalpha() for ch in stripped): score += 1 return score def generate_reply(model, prompt, user_text, stoi, itos, max_new_tokens, temperature, candidates=4): best_reply = "" best_score = float("-inf") for _ in range(candidates): reply = generate_once(model, prompt, stoi, itos, max_new_tokens, temperature) score = score_reply(reply, user_text) if score > best_score: best_score = score best_reply = reply return best_reply def train_steps(model, text, stoi, steps, batch_size, block_size, lr): if len(text) < block_size + 2: return data = encode_text(text, stoi) optimizer = torch.optim.AdamW(model.parameters(), lr=lr) model.train() for _ in range(steps): xb, yb = make_batch(data, batch_size, block_size, "cpu") _, loss = model(xb, yb) optimizer.zero_grad(set_to_none=True) loss.backward() optimizer.step() def save_model(model_path, model, config, stoi, itos): model_path.parent.mkdir(parents=True, exist_ok=True) torch.save( { "model": model.state_dict(), "config": config, "stoi": stoi, "itos": itos, }, model_path, ) def main(): parser = argparse.ArgumentParser() parser.add_argument("--data", default="data/chat_memory.txt") parser.add_argument("--seed-data", default="data/input.txt") parser.add_argument("--model", default="runs/chat_model.pt") parser.add_argument("--preset", choices=["tiny", "turbo", "small", "big", "large"], default="small") parser.add_argument("--steps-per-turn", type=int, default=8) parser.add_argument("--tokens", type=int, default=160) parser.add_argument("--temperature", type=float, default=1.1) parser.add_argument("--batch-size", type=int, default=4) parser.add_argument("--block-size", type=int, default=64) parser.add_argument("--n-embd", type=int, default=64) parser.add_argument("--n-head", type=int, default=2) parser.add_argument("--n-layer", type=int, default=1) parser.add_argument("--lr", type=float, default=3e-4) parser.add_argument("--no-self-train", action="store_true") args = parser.parse_args() presets = { "tiny": {"steps_per_turn": 4, "tokens": 120, "temperature": 1.0, "batch_size": 4, "block_size": 64, "n_embd": 64, "n_head": 2, "n_layer": 1, "lr": 3e-4}, "turbo": {"steps_per_turn": 2, "tokens": 100, "temperature": 1.2, "batch_size": 8, "block_size": 32, "n_embd": 64, "n_head": 4, "n_layer": 2, "lr": 1e-3}, "small": {"steps_per_turn": 8, "tokens": 160, "temperature": 1.1, "batch_size": 4, "block_size": 64, "n_embd": 96, "n_head": 2, "n_layer": 2, "lr": 2.5e-4}, "big": {"steps_per_turn": 12, "tokens": 180, "temperature": 1.0, "batch_size": 4, "block_size": 96, "n_embd": 192, "n_head": 4, "n_layer": 4, "lr": 2e-4}, "large": {"steps_per_turn": 16, "tokens": 220, "temperature": 0.95, "batch_size": 2, "block_size": 128, "n_embd": 256, "n_head": 8, "n_layer": 6, "lr": 1.5e-4}, } preset = presets[args.preset] if args.steps_per_turn == 8: args.steps_per_turn = preset["steps_per_turn"] if args.tokens == 160: args.tokens = preset["tokens"] if args.temperature == 1.1: args.temperature = preset["temperature"] if args.batch_size == 4: args.batch_size = preset["batch_size"] if args.block_size == 64: args.block_size = preset["block_size"] if args.n_embd == 64: args.n_embd = preset["n_embd"] if args.n_head == 2: args.n_head = preset["n_head"] if args.n_layer == 1: args.n_layer = preset["n_layer"] if args.lr == 3e-4: args.lr = preset["lr"] data_path = Path(args.data) seed_path = Path(args.seed_data) model_path = Path(args.model) if not data_path.exists(): seed = seed_path.read_text(encoding="utf-8") if seed_path.exists() else "" data_path.parent.mkdir(parents=True, exist_ok=True) data_path.write_text(seed + "\n", encoding="utf-8") text = data_path.read_text(encoding="utf-8") model, stoi, itos, config = load_or_create_model( model_path, text, args.block_size, args.n_embd, args.n_head, args.n_layer ) if not torch.cuda.is_available(): threads = os.cpu_count() or 4 torch.set_num_threads(threads) torch.set_num_interop_threads(1) torch.set_float32_matmul_precision("high") print(f"CPU optimization: using {threads} threads") print("Tiny chat. Type /quit to exit.") print("It uses retrieved examples and can keep training after each turn.") while True: try: user = input("\nyou> ").strip() if user.lower() in {"/quit", "quit", "exit"}: save_model(model_path, model, config, stoi, itos) print(f"saved {model_path}") break if user.startswith("/teach "): lesson = user[len("/teach ") :].strip() data_path.write_text(data_path.read_text(encoding="utf-8") + f"\nTEACHER: {lesson}\n", encoding="utf-8") limit_memory_file(data_path) text = data_path.read_text(encoding="utf-8") print("learning...") train_steps(model, text, stoi, args.steps_per_turn, args.batch_size, args.block_size, args.lr) save_model(model_path, model, config, stoi, itos) print("learned") continue memory_text = data_path.read_text(encoding="utf-8") context = retrieve_context(user, memory_text) if context: prompt = f"{context}\n\nUSER: {user}\nAI:" else: prompt = f"\nUSER: {user}\nAI:" reply = generate_reply(model, prompt, user, stoi, itos, args.tokens, args.temperature).strip() if not reply: reply = "..." reply = reply.replace("USER:", "").replace("AI:", "").strip() print(f"ai> {reply}") addition = f"\nUSER: {user}\nAI: {reply}\n" data_path.write_text(memory_text + addition, encoding="utf-8") limit_memory_file(data_path) text = data_path.read_text(encoding="utf-8") if not args.no_self_train: print("learning from this turn...") train_steps(model, text, stoi, args.steps_per_turn, args.batch_size, args.block_size, args.lr) save_model(model_path, model, config, stoi, itos) print("saved") except Exception as exc: print(f"error: {exc}") save_model(model_path, model, config, stoi, itos) print("model saved after error, continue or /quit") if __name__ == "__main__": main()