# app.py
import gradio as gr
import torch
import sentencepiece as spm
from karModel import DictionaryTransformer, device
from inference import generate_definition, encode, decode, sp

# Load model
vocab_size = sp.get_piece_size()
model = DictionaryTransformer(vocab_size).to(device)
model.load_state_dict(torch.load("cam/dictionary_transformer.pt", map_location=device))
model.eval()

def predict(word):
    if not word.endswith(":"):
        word += ":"
    return generate_definition(word)

demo = gr.Interface(
    fn=predict,
    inputs=gr.Textbox(label="Enter a word"),
    outputs=gr.Textbox(label="Generated Definition"),
    title="Dictionary Transformer",
    description="A small Transformer that generates dictionary-like definitions."
)

if __name__ == "__main__":
    demo.launch()