import gradio as gr
import torch
from transformers import GPT2LMHeadModel, GPT2Tokenizer

# Load the model and tokenizer
model_name = "abacaj/Replit-v2-CodeInstruct-3B-ggml"
model = GPT2LMHeadModel.from_pretrained(model_name)
tokenizer = GPT2Tokenizer.from_pretrained(model_name)

# Define the code generation function
def code_generation(code):
    inputs = tokenizer.encode(code, return_tensors="pt")
    outputs = model.generate(inputs, max_length=100, num_return_sequences=1)
    generated_code = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return generated_code

# Create the Gradio interface
iface = gr.Interface(
    fn=code_generation,
    inputs=gr.inputs.Textbox(lines=10, label="Enter your code"),
    outputs=gr.outputs.Textbox(label="Generated code"),
    title="Gardio App",
    description="An app that generates code based on user input.",
    examples=[
        ["Example input code snippet"],
        ["Another example input code snippet"],
    ],
    allow_screenshot=True
)

# Launch the interface
iface.launch()