import gradio as gr
# Check if CUDA (GPU) is available
import torch
from transformers import T5ForConditionalGeneration, PreTrainedTokenizerFast

# Define the path to the checkpoint directory
checkpoint_dir = "onlysainaa/cyrillic_to_script-t5-model"

# Load the model
model = T5ForConditionalGeneration.from_pretrained(checkpoint_dir)
model.eval()
# Load the tokenizer using PreTrainedTokenizerFast
tokenizer = PreTrainedTokenizerFast.from_pretrained(checkpoint_dir)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Move the model to the same device (GPU or CPU)
model.to(device)

# Function to perform translation using the model
def translate_text(input_text):
    # Tokenize the input text
    inputs = tokenizer(input_text, return_tensors="pt")
    
    # Move the input tensors to the same device as the model
    inputs = {k: v.to(device) for k, v in inputs.items() if k in ['input_ids', 'attention_mask']}
    
    # Generate translation
    outputs = model.generate(**inputs)
    
    # Decode the output to human-readable text
    translated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    return translated_text

# Create a Gradio interface
gr_interface = gr.Interface(
    fn=translate_text, 
    inputs="text", 
    outputs="text",
    title="Mongolian Cyrillic to Mongolian Script Model",
    description="Enter text in Mongolian Cyrillic"
)

# Launch the Gradio interface
gr_interface.launch()