import gradio as gr
import torch
import torch.nn.functional as F
import tiktoken
from model import GPT, GPTConfig  # We'll create this file next
from huggingface_hub import hf_hub_download

# Initialize model and load weights
def load_model():
    model = GPT(GPTConfig(
        block_size=128,
        vocab_size=50257,
        n_layer=12,
        n_head=12,
        n_embd=768
    ))
    
    # Download the model from the Hub
    model_path = hf_hub_download(
        repo_id="amlanr66/gpt2TransformerCustomeTrained",  # Replace with your model repository
        filename="best_model.pt"
    )
    
    checkpoint = torch.load(model_path, map_location=torch.device('cpu'))
    model.load_state_dict(checkpoint['model_state_dict'])
    model.eval()
    return model

# Text generation function
def generate_text(prompt, max_length=50, temperature=0.8, top_k=50):
    # Encode the prompt
    enc = tiktoken.get_encoding('gpt2')
    input_ids = torch.tensor(enc.encode(prompt)).unsqueeze(0)  # Add batch dimension
    
    # Generate text
    with torch.no_grad():
        for _ in range(max_length):
            # Get the last block_size tokens
            input_ids_cond = input_ids[:, -model.config.block_size:]
            
            # Forward pass
            logits = model(input_ids_cond)[0]
            logits = logits[:, -1, :] / temperature
            
            # Apply top-k sampling
            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = float('-inf')
            
            # Sample from the distribution
            probs = F.softmax(logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)
            
            # Append the new token
            input_ids = torch.cat((input_ids, next_token), dim=1)
    
    # Decode the generated text
    generated_text = enc.decode(input_ids[0].tolist())
    return generated_text

# Load the model globally
model = load_model()

# Create Gradio interface
def predict(text, length, temp):
    try:
        output = generate_text(
            prompt=text,
            max_length=int(length),
            temperature=float(temp)
        )
        return output
    except Exception as e:
        return f"Error: {str(e)}"

# Define the interface
iface = gr.Interface(
    fn=predict,
    inputs=[
        gr.Textbox(label="Input Text", placeholder="Enter your prompt here..."),
        gr.Slider(minimum=10, maximum=200, value=50, step=1, label="Maximum Length"),
        gr.Slider(minimum=0.1, maximum=2.0, value=0.8, step=0.1, label="Temperature")
    ],
    outputs=gr.Textbox(label="Generated Text"),
    title="GPT Text Generation",
    description="Enter a prompt and the model will continue the text.",
    examples=[
        ["The quick brown fox", 50, 0.8],
        ["Once upon a time", 100, 0.9],
        ["In a world where", 75, 0.7]
    ]
)

# Launch the app
if __name__ == "__main__":
    iface.launch()