chat_train.py

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