from difflib import Differ

import gradio as gr
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

model = AutoModelForSeq2SeqLM.from_pretrained("Hnabil/t5-address-standardizer")
tokenizer = AutoTokenizer.from_pretrained("Hnabil/t5-address-standardizer")


def standardizer_adress(adress):
    inputs = tokenizer(adress, return_tensors="pt")
    outputs = model.generate(**inputs, max_length=100)
    std_adress = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
    return std_adress


def diff_texts(adress):
    std_adress = standardizer_adress(adress)
    words1 = adress.split()
    words2 = std_adress.split()
    d = Differ()

    return [
        (token[2:], "Non-Standard" if token[0] != " " else None)
        for token in d.compare(words1, words2)
        if token[0] != "+" and token[0] != "?"
    ] + [
        ("→", "→")
    ] + [
        (token[2:], "Standard" if token[0] != " " else None)
        for token in d.compare(words1, words2)
        if token[0] != "-" and token[0] != "?"
    ]


examples = [
    ["940, north pennsylvania avneue, mason icty, iowa, 50401, us"],
    ["537, 6th st s, mason city, ia, 50401, us"]
]
color_map = {"Non-Standard": "LightSalmon", "Standard": "LightGreen", "→": "LightBlue"}

demo = gr.Interface(
    diff_texts,
    [
        gr.Textbox(
            label="Adress",
            info="Enter an adress to standardize",
            lines=3,
        )
    ],
    [
      gr.HighlightedText(
          label="Standardized Adress",
          show_legend=True,
      ).style(color_map=color_map),
    ],
    examples=examples,
    theme=gr.themes.Base()
)
if __name__ == "__main__":
    demo.launch(share=False)