from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import gradio as gr

# Load the model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("sagard21/python-code-explainer")
model = AutoModelForSeq2SeqLM.from_pretrained("sagard21/python-code-explainer")

def explain_code(python_code):
    # Tokenize the input code
    inputs = tokenizer(python_code, return_tensors="pt", truncation=True, max_length=512)
    
    # Generate explanation
    explanation_ids = model.generate(
        inputs["input_ids"],
        max_length=256,
        num_beams=5,
        early_stopping=True
    )
    
    # Decode and return the explanation
    explanation = tokenizer.decode(explanation_ids[0], skip_special_tokens=True)
    return explanation

# Create the Gradio interface
demo = gr.Interface(
    fn=explain_code,
    inputs=gr.Code(
        language="python",
        label="Enter Python Code",
        lines=10,
        placeholder="def hello_world():\n    print('Hello, world!')"
    ),
    outputs=gr.Textbox(
        label="Code Explanation",
        lines=5
    ),
    title="Python Code Explainer",
    description="🔍 Enter Python code and get a natural language explanation of what it does.",
    examples=[
        ["def add(a, b):\n    return a + b"],
        ["for i in range(5):\n    print(i)"],
        ["x = [i**2 for i in range(10) if i % 2 == 0]"]
    ]
)

# Launch the app
demo.launch()