import gradio as gr
import torch
from transformers import T5Tokenizer, T5ForConditionalGeneration, Trainer, TrainingArguments


def transliteration(source_word):
    #return "Hello " + name + "!!"
    #source_word = "Manik aai Ameela maqaamun vakikuran hushahalhaifi"
    source_word_str = source_word.lower()

    tokenizer = T5Tokenizer.from_pretrained("chopey/dvt5-base")
    model = T5ForConditionalGeneration.from_pretrained('chopey/model_t5_base')
    
    input_ids = tokenizer.encode(source_word_str, return_tensors="pt")
    output_ids = model.generate(input_ids, max_length=512)
    transliteration = tokenizer.decode(output_ids[0], skip_special_tokens=True)
    
    return transliteration

#test
iface = gr.Interface(fn=transliteration, inputs="text", outputs="text")
iface.launch()