import gradio as gr
from transformers import pipeline
fill_mask = pipeline('fill-mask',
                 model='davmel/ka_homonym_disambiguation_FM',
                 tokenizer='davmel/ka_homonym_disambiguation_FM')  # 1788?

masker = {0: 'თო', 1: 'დაბ', 2: 'კაფე'}

def getMasked(sent):
    homonym_forms = [
        'ბარით', 'ბარიანის', 'ბარზე', 'ბარები', 'ბარისთვისაც', 'ბარებია', 'ბარეთსაც', 'ბარაინში', 'ბარეთმაც',
        'ბარიანი',
        'ბარიცაა', 'ბარიდანა', 'ბარიც', 'ბარვაზედ', 'ბარელიც', 'ბარავდი', 'ბარისკაცები', 'ბარვაში', 'ბარვას',
        'ბარეთისა', 'ბარავოოო', 'ბარისებური', 'ბარებისთვისა', 'ბარისათვის', 'ბარზეა', 'ბარელი', 'ბარშივე',
        'ბარისაა',
        'ბარისკენ', 'ბარელებით', 'ბარადა', 'ბარვაა', 'ბარს', 'ბარიო', 'ბარავდეს', 'ბარებმა', 'ბარებით', 'ბარების',
        'ბარისგან', 'ბარშიო', 'ბარებისკენ', 'ბარიდანაა', 'ბარებინა', 'ბარებისა', 'ბარვისგან', 'ბარვისთან',
        'ბარავდნენ',
        'ბარადაც', 'ბარელამდე', 'ბარშიაო', 'ბარადაა', 'ბარითა', 'ბარელებიც', 'ბარელის', 'ბარელებს', 'ბარვის',
        'ბარელებმა', 'ბარელმა', 'ბარადო', 'ბარებისგან', 'ბარვისას', 'ბარებშია', 'ბარავენ', 'ბარვით', 'ბარიანს',
        'ბარელებში', 'ბარიცო', 'ბარავდა', 'ბარშია', 'ბარვაზე', 'ბარავთ', 'ბარო', 'ბარელისა', 'ბარაინი',
        'ბარელისათვის',
        'ბარელია', 'ბარებშიო', 'ბარის', 'ბარისთვისო', 'ბარებად', 'ბარშიგან', 'ბარებსაც', 'ბარებითდა', 'ბარებზე',
        'ბარია', 'ბარებიდან', 'ბარისად', 'ბარზეც', 'ბარშიც', 'ბარებითაა', 'ბარად', 'ბარიდანაც', 'ბარისმაგვარ',
        'ბარში',
        'ბარებისთვისაც', 'ბართანა', 'ბარშიცაა', 'ბარისვე', 'ბარისთვისა', 'ბარმაც', 'ბარელით', 'ბარვა', 'ბარელებთან',
        'ბარელებზე', 'ბარისაც', 'ბარისებურ', 'ბარელების', 'ბარიდან', 'ბარისი', 'ბარებმაც', 'ბარივით', 'ბარიანმა',
        'ბარელისთვის', 'ბარებს', 'ბარ', 'ბარებიც', 'ბარებზეც', 'ბარებისთვის', 'ბარელსაც', 'ბარელები', 'ბარისია',
        'ბარისპირით', 'ბარაამდეე', 'ბარამდეც', 'ბარი', 'ბარისაკენ', 'ბარელად', 'ბარელ', 'ბარამდე', 'ბარიანად',
        'ბარებისათვის', 'ბარიანია', 'ბარებში', 'ბართან', 'ბარეთთან', 'ბარისაგან', 'ბარებთაგან', 'ბარისპირის',
        'ბარზედა',
        'ბარვაო', 'ბარებშიდაა', 'ბარებსადა', 'ბარელამდეც', 'ბარებითა', 'ბარისა', 'ბარებშიც', 'ბარიდანვე', 'ბარისო',
        'ბარიდანნ', 'ბარელს', 'ბარელთა', 'ბართა', 'ბართანაც', 'ბარავო', 'ბარისთვის', 'ბარელზეა', 'ბარებსა',
        'ბარისაო',
        'ბარებიზე', 'ბარავს', 'ბარისმაგვარი', 'ბარებთან', 'ბარებითო', 'ბარსა', 'ბარული', 'ბარსაც', 'საბარო',
        'ბარებითო',
        'ბარშიაც', 'ბარობა', 'ბარისას',
    ]
    new_sent = []
    wasnt = True
    for word in sent.split(" "):
        if word in homonym_forms and wasnt:
            new_sent.append("[MASK]")
            # wasnt = False
        else:
            new_sent.append(word)
    return " ".join(new_sent)

def predict(text):
    # Mask some words in the input text
    text_with_mask = getMasked(text)
    # Use the model to predict the masked words
    predictions = fill_mask(text_with_mask)
    # Extract the predicted word and score
    context = predictions[0]["token_str"]
    if context == masker[0]:
        return "Homonym is used as a \"Shovel\""
    elif context == masker[1]:
        return "Homonym is used as a \"Lowland\""
    else:
        return "Homonym is used as a \"Cafe\""

# Create the Gradio interface
iface = gr.Interface(fn=predict, inputs="text", outputs="text", title="Homonym disambiguation in Georgian",
                     description="Enter a sentence with the homonym \"ბარი\" (for the current purposes, please include the homonym once in the sentence).")
iface.launch(share=True)