import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# ✅ Use AutoTokenizer and AutoModel to avoid tokenizer mismatch
tokenizer = AutoTokenizer.from_pretrained("Salesforce/codet5-base")
model = AutoModelForSeq2SeqLM.from_pretrained("Salesforce/codet5-base")

# Function to explain code
def explain_code(code_snippet):
    if not code_snippet.strip():
        return "❗ Please enter some code."
    
    input_text = f"summarize: {code_snippet.strip()}"
    input_ids = tokenizer.encode(input_text, return_tensors="pt", truncation=True, max_length=512)
    outputs = model.generate(input_ids, max_length=150, num_beams=4, early_stopping=True)
    explanation = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    return explanation

# Gradio Interface
demo = gr.Interface(
    fn=explain_code,
    inputs=gr.Textbox(lines=15, label="Paste your code here"),
    outputs=gr.Textbox(label="Explanation"),
    title="🧠 Code Explainer using CodeT5",
    description="Paste your code and get a plain English explanation using Salesforce's CodeT5.",
    theme="default"
)

demo.launch()