chat.py

"""GPT model chat interface.

This module provides an interactive chat interface for a trained GPT model,
allowing users to select a model from the models directory and generate
responses to prompts.
"""

from pathlib import Path

import sentencepiece as spm
import torch
import torch.nn.functional as F

from train_gpt import GPT, Hyperparameters


def select_model() -> str | None:
    """Select a model file from the 'models' directory.

    Returns:
        Path to selected model file as string, or None if no models found.

    """
    models_dir = Path("models")

    # Check if models directory exists
    if not models_dir.exists():
        print("Models directory 'models' not found!")
        return None

    # Search for .pt files in the models directory
    models = [f for f in models_dir.glob("*.pt")]

    if not models:
        print("No model files (.pt) found in 'models' directory!")
        return None

    print("\nAvailable models:")
    for i, model_path in enumerate(models):
        print(f"{i + 1}: {model_path.name}")  # Show only filename, not full path

    try:
        idx = int(input("\nSelect model number for chat: ")) - 1
        if 0 <= idx < len(models):
            return str(models[idx])  # Return full path as string
        print("Invalid model number!")
        return None
    except ValueError:
        print("Please enter a valid number!")
        return None


@torch.no_grad()
def generate(model: GPT, tokenizer: spm.SentencePieceProcessor, prompt: str, max_new_tokens: int = 30) -> str:
    """Generate text continuation from a prompt using the model.

    Args:
        model: The GPT model instance.
        tokenizer: SentencePiece tokenizer for encoding/decoding.
        prompt: Input text prompt.
        max_new_tokens: Maximum number of new tokens to generate.

    Returns:
        Generated text including the original prompt.

    """
    model.eval()

    # Encode the input prompt
    tokens = tokenizer.encode(prompt)
    x = torch.tensor([tokens], dtype=torch.long, device="cuda")

    for _ in range(max_new_tokens):
        # Manual forward pass to get logits for generation
        # The original forward returns CrossEntropy loss, so we replicate
        # the forward logic here to extract logits

        # Token embeddings
        x_emb = model.tok_emb(x)
        x_emb = F.rms_norm(x_emb, (x_emb.size(-1),))
        x0 = x_emb
        skips = []

        # Encoder layers
        for i in range(model.num_encoder_layers):
            x_emb = model.blocks[i](x_emb, x0)
            skips.append(x_emb)

        # Decoder layers with skip connections
        for i in range(model.num_decoder_layers):
            if skips:
                x_emb = x_emb + model.skip_weights[i].to(dtype=x_emb.dtype)[None, None, :] * skips.pop()
            x_emb = model.blocks[model.num_encoder_layers + i](x_emb, x0)

        # Take only the last token
        x_final = model.final_norm(x_emb[:, -1, :])

        # Get logits
        if model.tie_embeddings:
            logits = F.linear(x_final, model.tok_emb.weight)
        else:
            logits = model.lm_head(x_final)

        logits = model.logit_softcap * torch.tanh(logits / model.logit_softcap)

        # Sample next token
        probs = F.softmax(logits, dim=-1)
        next_token = torch.multinomial(probs, num_samples=1)

        # Append to sequence
        x = torch.cat((x, next_token), dim=1)

        # Stop if EOS token is generated
        if next_token.item() == tokenizer.eos_id():
            break

    return tokenizer.decode(x[0].tolist())


if __name__ == "__main__":
    # Initialize device
    device = "cuda"
    hp = Hyperparameters

    # Select and load model
    model_path = select_model()
    if model_path:
        # Instantiate model with hyperparameters
        model = GPT(
            vocab_size=hp.vocab_size,
            num_layers=hp.num_layers,
            model_dim=hp.model_dim,
            num_heads=hp.num_heads,
            num_kv_heads=hp.num_kv_heads,
            mlp_mult=hp.mlp_mult,
            tie_embeddings=hp.tie_embeddings,
            tied_embed_init_std=hp.tied_embed_init_std,
            logit_softcap=hp.logit_softcap,
            rope_base=hp.rope_base,
            qk_gain_init=hp.qk_gain_init,
        ).to(device)

        # Load model weights
        state_dict = torch.load(model_path, map_location=device)
        model.load_state_dict(state_dict)
        print(f"Model {model_path} successfully loaded!")

        # Initialize tokenizer
        tokenizer = spm.SentencePieceProcessor(model_file=hp.tokenizer_path)

        # Start chat interface
        print("\n--- Chat Interface (type 'exit' to quit) ---")
        while True:
            user_input = input("\nYou: ")
            if user_input.lower() == "exit":
                break

            try:
                output = generate(model, tokenizer, user_input)
                print(f"AI:  {output}")
            except Exception as e:
                print(f"Generation error: {e}")