chat.py

#!/usr/bin/env python3
"""
LLM Chat Interface

Usage:
    python chat.py
"""

import torch
import pickle
import argparse
import os


def check_requirements():
    """Check if required files exist"""
    required_files = ["vocab.txt", "llm_model.pkl"]
    missing = [f for f in required_files if not os.path.exists(f)]

    if missing:
        print("Missing required files:")
        for f in missing:
            print(f"  - {f}")
        print("\nPlease ensure all required files are present.")
        return False
    return True


def main():
    if not check_requirements():
        return

    # Configuration


device = (
    "cuda"
    if torch.cuda.is_available()
    else "mps"
    if torch.backends.mps.is_available()
    else "cpu"
)
block_size = 128

print(f"Using device: {device}")

# Load vocabulary
with open("vocab.txt", "r", encoding="utf-8") as f:
    text = f.read()
    chars = sorted(list(set(text)))

vocab_size = len(chars)

# Create encoding/decoding maps
string_to_int = {ch: i for i, ch in enumerate(chars)}
int_to_string = {i: ch for i, ch in enumerate(chars)}
encode = lambda s: [string_to_int[c] for c in s]
decode = lambda l: "".join([int_to_string[i] for i in l])

# Load trained model
with open("llm_model.pkl", "rb") as f:
    model = pickle.load(f)
    model = model.to(device)

print("Model loaded successfully!")


# Interactive chat function
def interactive_chat():
    print("\nInteractive Chat Mode")
    print("Type 'quit' to exit")

    while True:
        try:
            prompt = input("\nYou: ")
            if prompt.lower() == "quit":
                break

            if prompt.strip() == "":
                continue

            context = torch.tensor(encode(prompt), dtype=torch.long, device=device)
            context = context.unsqueeze(0)  # Add batch dimension

            # Generate response
            max_new_tokens = 100
            generated_tensor = model.generate(context, max_new_tokens)
            generated_chars = decode(generated_tensor[0].tolist())

            # Extract only the generated part (after the prompt)
            generated_text = generated_chars[len(prompt) :]

            print(f"LLM: {generated_text}")

        except KeyboardInterrupt:
            print("\nExiting...")
            break
        except Exception as e:
            print(f"Error: {e}")


def test_generation():
    prompt = "Hello! Can you see me?"
    context = torch.tensor(encode(prompt), dtype=torch.long, device=device)
    context = context.unsqueeze(0)
    generated_tensor = model.generate(context, 100)
    generated_chars = decode(generated_tensor[0].tolist())
    generated_text = generated_chars[len(prompt) :]
    print(f"Test generation:\nPrompt: {prompt}\nGenerated: {generated_text}")


def main():
    parser = argparse.ArgumentParser(description="LLM Chat Interface")
    parser.add_argument("--test", action="store_true", help="Test generation")
    args = parser.parse_args()

    if args.test:
        test_generation()
    else:
        interactive_chat()


if __name__ == "__main__":
    main()