import gradio as gr
import tensorflow as tf
from transformers import AutoTokenizer, TFAutoModelForSeq2SeqLM

path = 'tf_model/'
model_checkpoint = "Helsinki-NLP/opus-mt-en-hi"
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
model = TFAutoModelForSeq2SeqLM.from_pretrained(path)

title = 'Text Translation(English to Hindi)'
def process_input(text):
    # Tokenize the input text using the tokenizer and convert to NumPy arrays
    tokenized = tokenizer([text], return_tensors='np')
    # Generate output sequences using the pre-trained model
    out = model.generate(**tokenized, max_length=128)
    # Switch the tokenizer to target mode
    with tokenizer.as_target_tokenizer():
        # Decode the generated output sequence, skipping special tokens
        result = tokenizer.decode(out[0], skip_special_tokens=True)
    return result

# Example input text for the GUI
examples = ['If you have the time, come along with me.', 'I can come if you want.', 'Tom was at home alone.', 'Wow!','How rude of you!',"What's in your hand?"]

# Create a Gradio Interface for the model
model_gui = gr.Interface(
    process_input,                 # Function for processing input and generating output
    gr.Textbox(lines=3, label="English"),  # Textbox for entering English text
    gr.Textbox(lines=3, label="Hindi"),    # Textbox for displaying translated Hindi text
    title=title,                   # Set the title of the GUI
    examples=examples              # Provide example input text for the GUI
)

# Launch the Gradio GUI with sharing enabled
model_gui.launch()