Upload chat_train.py with huggingface_hub

chat_train.py

@@ -0,0 +1,282 @@
+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()