import gradio as gr
from transformers import pipeline
import torch

# Define the prompt template
MAGICODER_PROMPT = """You are an exceptionally intelligent coding assistant that consistently delivers accurate and reliable responses to user instructions.

@@ Instruction
{instruction}

@@ Response
"""

# Load the Magicoder model
generator = pipeline(
    model="ise-uiuc/Magicoder-S-DS-6.7B",
    task="text-generation",
    torch_dtype=torch.bfloat16,
    device_map="auto",
)

# Define the function to use with Gradio
def generate_response(instruction):
    prompt = MAGICODER_PROMPT.format(instruction=instruction)
    result = generator(prompt, max_length=2048, num_return_sequences=1, temperature=0.0)
    return result[0]["generated_text"]

# Create the Gradio interface
demo = gr.Interface(fn=generate_response, inputs="text", outputs="text")

# Launch the interface
demo.launch(share=True)