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

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

# Translation function
def translate_latin_to_italian(latin_text):
    # Split input text into sentences while preserving line breaks
    sentences = re.split(r'(?<=[.!?]) +', latin_text.strip())
    
    translated_sentences = []
    
    for sentence in sentences:
        # Make the first letter lowercase if the sentence is not empty
        if sentence:
            sentence = sentence[0].lower() + sentence[1:]
        inputs = tokenizer(sentence, 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)
        translated_sentences.append(translation[0])
    
    # Reassemble the translated sentences and keep original line breaks
    translated_text = ' '.join(translated_sentences)
    return translated_text

# 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()