import gradio as gr
import torch
import torch.nn as nn

# Define your custom model class
class BigramLanguageModel(nn.Module):
    def __init__(self):
        super().__init__()
        # Example layers (adjust as needed for your model)
        self.token_embedding_table = nn.Embedding(61, 64)
        self.position_embedding_table = nn.Embedding(32, 64)
        self.blocks = nn.Sequential(*[nn.Linear(64, 64) for _ in range(4)])
        self.ln_f = nn.LayerNorm(64)
        self.lm_head = nn.Linear(64, 61)

    def forward(self, idx):
        # Implement the forward pass
        pass

    def generate(self, idx, max_new_tokens=250):
        # Implement the generate method
        pass

# Load your model
def load_model():
    model = BigramLanguageModel()
    model_url = "https://huggingface.co/yoonusajwardapiit/triptuner/resolve/main/pytorch_model.bin"
    model_weights = torch.hub.load_state_dict_from_url(model_url, map_location=torch.device('cpu'), weights_only=True)
    model.load_state_dict(model_weights)
    model.eval()
    return model

model = load_model()

# Define encode and decode functions
chars = sorted(list(set("your_training_text_here")))  # Replace with the character set used in training
stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for i, ch in enumerate(chars)}
encode = lambda s: [stoi[c] for c in s]
decode = lambda l: ''.join([itos[i] for i in l])

# Function to generate text using the model
def generate_text(prompt):
    context = torch.tensor([encode(prompt)], dtype=torch.long)
    with torch.no_grad():
        generated = model.generate(context, max_new_tokens=250)  # Adjust as needed
    return decode(generated[0].tolist())

# Create a Gradio interface
interface = gr.Interface(
    fn=generate_text,
    inputs=gr.Textbox(lines=2, placeholder="Enter a location or prompt..."),
    outputs="text",
    title="Triptuner Model",
    description="Generate itineraries for locations in Sri Lanka's Central Province."
)

# Launch the interface
interface.launch()