import gradio as gr
from transformers import MarianMTModel, MarianTokenizer
import torch

# Load the model and tokenizer from the Hub
model_name = "Dddixyy/latin-italian-translatorV5"
tokenizer = MarianTokenizer.from_pretrained(model_name)
model = MarianMTModel.from_pretrained(model_name)

# Translation function
def translate_latin_to_italian(latin_text):
    # Make the first letter lowercase if the input is not empty
    if latin_text:
        latin_text = latin_text[0].lower() + latin_text[1:]
    inputs = tokenizer(latin_text, return_tensors="pt", padding=True, truncation=True)
    with torch.no_grad():
        generated_ids = model.generate(inputs["input_ids"])
    translation = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
    return translation[0]

# Define the Gradio interface
interface = gr.Interface(
    fn=translate_latin_to_italian,
    inputs="text",
    outputs="text",
    title="Latin to Italian Translator",
    description="Translate Latin sentences to Italian using a fine-tuned MarianMT model.",
    examples=[
        ["Amor vincit omnia."],
        ["Veni, vidi, vici."],
        ["Carpe diem."],
        ["Alea iacta est."]
    ]
)

# Launch the app
if __name__ == "__main__":
    interface.launch()