import gradio as gr
import os
import string
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

hf_token = os.environ.get('HF_TOKEN', None)

tokenizer = AutoTokenizer.from_pretrained("hajili/mt5-spell-correction-azerbaijani", token=hf_token)
model = AutoModelForSeq2SeqLM.from_pretrained("hajili/mt5-spell-correction-azerbaijani", token=hf_token)

tokenizer_new = AutoTokenizer.from_pretrained("hajili/spell-correction-azerbaijani-v2", token=hf_token)
model_new = AutoModelForSeq2SeqLM.from_pretrained("hajili/spell-correction-azerbaijani-v2", token=hf_token)

def correct(sentence):
    sentence = sentence.lower().translate(str.maketrans('','', string.punctuation)).strip()
    
    batch = tokenizer(sentence, return_tensors="pt")
    generated_ids = model.generate(batch["input_ids"])
    output = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)

    batch_new = tokenizer_new(sentence, return_tensors="pt")
    generated_ids_new = model_new.generate(batch_new["input_ids"])
    output_new = tokenizer_new.batch_decode(generated_ids_new, skip_special_tokens=True)

    return output[0], output_new[0]

input_component = gr.Textbox(label = "Giriş cümləsini yazın")
output_component = gr.Textbox(label = "Modelin cavabına əsasən verilən cümlənin düzgün yazılışı:")
output_component_new = gr.Textbox(label = "Yeni modelin cavabına əsasən verilən cümlənin düzgün yazılışı:")
examples = [["inweallah duzeler qardawim ureyini cox sixma her wey yaxwi olacax"], ["veziyet hec de urek acan deyil hami xeberi gozleyir ki gorek ne olacaq imkanin olsa sen de zeng eleyib bri xeber sorusharsan xahis edirem anan qurban gozderine"]]
description = "- Bu model Azərbaycan dilində yazı səhvlərini düzəltmək üçün dizayn edilib.\n- Modeli test etmək üçün giriş cümləsini daxil edin."
demo = gr.Interface(correct, inputs = input_component, outputs=[output_component, output_component_new], examples=examples, title = "Azərbaycan dilində yazı səhvlərinin düzəlişi üçün dil modeli", description=description)

demo.launch()