import gradio as gr
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

# Load the model and tokenizer
model_name = "SnypzZz/Llama2-13b-Language-translate"
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Function to handle translation
def translate_text(text, target_language):
    # Tokenize the input text
    inputs = tokenizer.encode(text, return_tensors="pt")
    
    # Perform translation
    outputs = model.generate(inputs)
    
    # Decode the translated output
    translated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return translated_text

# Gradio interface
def gradio_app():
    # Create a Gradio interface
    iface = gr.Interface(
        fn=translate_text, 
        inputs=[
            gr.inputs.Textbox(label="Input Text"), 
            gr.inputs.Dropdown(choices=["fr", "es", "de", "zh"], label="Target Language")  # Modify languages accordingly
        ],
        outputs=gr.outputs.Textbox(label="Translated Text"),
        title="Language Translator",
        description="A translation app using SnypzZz/Llama2-13b-Language-translate"
    )
    
    # Launch the interface
    iface.launch()

if __name__ == "__main__":
    gradio_app()