import torch
from model_transformer import TransformerLM
from tokenizers import Tokenizer
import gradio as gr

# Load tokenizer
tok = Tokenizer.from_file("tokenizer.json")

def encode(text):
    return tok.encode(text).ids

def decode(ids):
    return tok.decode(ids)

# Load model
vocab_size = tok.get_vocab_size()
model = TransformerLM(vocab_size)
model.load_state_dict(torch.load("model10M.pt", map_location="cpu"))
model.eval()

# Text generation
def generate(prompt, max_len=100):
    ids = encode(prompt)
    ids = torch.tensor([ids], dtype=torch.long)

    for _ in range(max_len):
        with torch.no_grad():
            logits = model(ids)
        next_id = torch.argmax(logits[0, -1]).item()
        ids = torch.cat([ids, torch.tensor([[next_id]])], dim=1)

    output = decode(ids[0].tolist())
    return output

# Gradio UI
demo = gr.Interface(
    fn=generate,
    inputs=gr.Textbox(lines=2, placeholder="Ask something..."),
    outputs="text",
    title="ChudAI (Sandesh Edition)",
    description="Your custom 10M Transformer AI running on HuggingFace."
)

demo.launch()