import gradio as gr
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

# Load the model and tokenizer
model_name = "facebook/codegen-350M-mono"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

# Function to translate Python code to JavaScript
def translate_code(python_code):
    inputs = tokenizer(python_code, return_tensors="pt", max_length=512, truncation=True)
    outputs = model.generate(**inputs, max_length=512, num_beams=5, early_stopping=True)
    translated_code = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return translated_code

# Gradio interface
interface = gr.Interface(
    fn=translate_code, 
    inputs="text", 
    outputs="text",
    title="Python to JavaScript Code Translator",
    description="Paste your Python code, and it will be translated into JavaScript.",
)

# Launch the app
interface.launch()