import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# Load CodeGen model and tokenizer
model_name = "Salesforce/codegen-2B-mono"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)

def generate_response(input_text, max_length=250, temperature=0.7, top_p=0.9, top_k=50):
    """
    Generate response using the CodeGen model based on user input and selected parameters.
    """
    try:
        # Encode input and prepare input tensor
        inputs = tokenizer(input_text, return_tensors="pt").to(device)
        
        # Generate text based on model output
        outputs = model.generate(
            inputs.input_ids,
            max_length=max_length,
            temperature=temperature,
            top_p=top_p,
            top_k=top_k,
            do_sample=True,
            num_return_sequences=1,
            no_repeat_ngram_size=2
        )
        
        # Decode and return the generated text
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        return response
    
    except Exception as e:
        return f"Error generating response: {str(e)}"

# Create Gradio interface
with gr.Blocks() as codegen_app:
    gr.Markdown("# CodeGen-powered Text Generation")

    # Input box for user prompt
    input_text = gr.Textbox(
        label="Input Text",
        placeholder="Type your question or prompt here",
        lines=3
    )

    # Sliders for customization
    max_length = gr.Slider(
        label="Max Length",
        minimum=50,
        maximum=1024,
        step=10,
        value=250
    )
    temperature = gr.Slider(
        label="Temperature",
        minimum=0.1,
        maximum=1.0,
        step=0.1,
        value=0.7
    )
    top_p = gr.Slider(
        label="Top-p (Nucleus Sampling)",
        minimum=0.1,
        maximum=1.0,
        step=0.1,
        value=0.9
    )
    top_k = gr.Slider(
        label="Top-k (Sampling Limit)",
        minimum=0,
        maximum=100,
        step=5,
        value=50
    )

    # Output box
    output_text = gr.Textbox(
        label="Generated Response",
        placeholder="The model's response will appear here",
        lines=15
    )

    # Generate button
    generate_button = gr.Button("Generate Response")
    generate_button.click(
        fn=generate_response,
        inputs=[input_text, max_length, temperature, top_p, top_k],
        outputs=output_text
    )

# Launch the app
codegen_app.launch()