import gradio as gr
import torch
from transformers import MarianMTModel, MarianTokenizer
from optimum.intel import IncQuantizer

# Load and optimize the model (quantization)
model_name = "Dddixyy/latin-italian-translator"

# Load the quantized model if available or use a regular model (quantization shown as an example)
try:
    # Attempt to load a quantized version if it's available
    quantizer = IncQuantizer.from_pretrained(model_name)
    model = quantizer.quantize()
    print("Quantized model loaded.")
except Exception as e:
    print(f"Error loading quantized model: {e}")
    model = MarianMTModel.from_pretrained(model_name)

# Load tokenizer
tokenizer = MarianTokenizer.from_pretrained(model_name)

# Translation function
def translate_latin_to_italian(latin_text):
    # Truncate input to 512 tokens to avoid overload (adjust as necessary)
    inputs = tokenizer(latin_text, return_tensors="pt", padding=True, truncation=True, max_length=512)
    
    with torch.no_grad():
        generated_ids = model.generate(inputs["input_ids"])

    # Decode the generated ids into a readable translation
    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(server_name="0.0.0.0", server_port=7860)