from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

# Path to your model's checkpoints
model_checkpoint_path = "barbaroo/nllb_200_600M_fo_en"

# Load the tokenizer and model
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint_path)
model = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint_path)

# Function to perform translation
def translate(text, model, tokenizer):
    # Encode the text
    inputs = tokenizer.encode(text, return_tensors="pt")

    # Generate translation outputs
    outputs = model.generate(inputs, max_length=80, num_beams=4, early_stopping=True)

    # Decode and return the translated text
    translated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return translated_text

# Sample translation (You can comment this out when deploying to Gradio)
faroese_text = "Granskingarvirksemi fevnir bæði um føroyskar og altjóða granskingarverkætlanir."
translated_text = translate(faroese_text, model, tokenizer)
print(translated_text)

# Gradio Interface setup
# Ensure Gradio is installed
!pip install gradio

# Importing Gradio
import gradio as gr

# Define the Gradio interface
def gradio_translate(text):
    return translate(text, model, tokenizer)

iface = gr.Interface(fn=gradio_translate, 
                     inputs="text", 
                     outputs="text", 
                     title="Faroese to English Translator",
                     description="Translate Faroese text to English using a state-of-the-art model.")

# Launch the interface
iface.launch()