import os
import numpy as np
import onnxruntime as ort
from transformers import AutoTokenizer
import gradio as gr

# Load the ONNX model and tokenizer
model_path = "model.onnx"
translation_session = ort.InferenceSession(model_path)
translation_tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-fr")

def translate_text(input_text):
    # Tokenize input text
    tokenized_input = translation_tokenizer(
        input_text, return_tensors="np", padding=True, truncation=True, max_length=512
    )
    
    input_ids = tokenized_input["input_ids"].astype(np.int64)
    attention_mask = tokenized_input["attention_mask"].astype(np.int64)

    # Run inference with the ONNX model
    outputs = translation_session.run(
        None,
        {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
        }
    )

    # Decode the output tokens
    translated_tokens = np.argmax(outputs[0], axis=-1)
    translated_text = translation_tokenizer.decode(translated_tokens[0], skip_special_tokens=True)

    return translated_text

# Create a Gradio interface
interface = gr.Interface(
    fn=translate_text,
    inputs="text",
    outputs="text",
    title="Frenchizer Translation Model",
    description="Translate text from English to French using an ONNX model."
)

# Launch the Gradio app
interface.launch()