import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM

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

# Function to translate Python code to JavaScript
def translate_code(python_code):
    inputs = tokenizer(python_code, return_tensors="pt")
    outputs = model.generate(inputs["input_ids"], max_length=150, num_return_sequences=1)
    translated_code = tokenizer.decode(outputs[0], skip_special_tokens=True, clean_up_tokenization_spaces=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()