import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_NAME = "oskaralf/model_merged" 
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype="auto", device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

def generate_response(prompt, max_length=128, temperature=0.7, top_p=0.9):
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(
        **inputs,
        max_length=max_length,
        temperature=temperature,
        top_p=top_p,
        pad_token_id=tokenizer.eos_token_id
    )
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

def interactive_app():
    with gr.Blocks() as app:
        gr.Markdown("# Coding Task Generator")
        gr.Markdown("Generate coding tasks by entering a prompt below.")
        
        prompt = gr.Textbox(label="Enter your prompt:", placeholder="e.g., Create a Python task involving recursion.")
        max_length = gr.Slider(label="Max Length", minimum=16, maximum=512, value=128, step=16)
        temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=1.0, value=0.7, step=0.1)
        top_p = gr.Slider(label="Top-p Sampling", minimum=0.1, maximum=1.0, value=0.9, step=0.1)
        generate_button = gr.Button("Generate Task")
        
        output = gr.Textbox(label="Generated Task", lines=10)
        
        generate_button.click(
            generate_response, 
            inputs=[prompt, max_length, temperature, top_p],
            outputs=output
        )
    
    return app

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