| 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() |
|
|
