import gradio as gr
from transformers import pipeline

title = "Fold: Contextual Tag Recommendation System"
description = "powered by bart-large-mnli, made by @abhisheky127"

classifier = pipeline("zero-shot-classification",
                      model="facebook/bart-large-mnli")

#define a function to process your input and output
def zero_shot(doc, candidates):
    given_labels = candidates.split(", ")
    dictionary = classifier(doc, given_labels)
    labels = dictionary['labels']
    scores = dictionary['scores']
    return dict(zip(labels, scores))

#define a function to preprocess transaction query
def preprocess(transaction):
    pattern = r'([A-Za-z0-9\s]+)(?:/| |$)'
    match = re.search(pattern, transaction)
    if match:
        return match.group(1).strip()
    return None


#create input and output objects
#input object1
input1 = gr.Textbox(label="Text")

#input object 2
input2 = gr.Textbox(label="Labels")

#output object
output = gr.Label(label="Output")

#example object
transactions_and_tags = [
    ["MPS/TRUFFLES /202303261700/034587/Bangalore", "Medical,Food,Shopping,Subscription,Travel, Miscellaneous"],
    ["MPS/TACO BELL /202304012247/108300/BANGALORE", "Medical,Food,Shopping,Subscription,Travel, Miscellaneous"],
    ["POS XXXXXXXXXXXX0001 APOLLO PHARMACY", "Medical,Food,Shopping,Subscription,Travel, Miscellaneous"],
    ["BIL/ONL/000471093694/1MG Techno/X7ZRUSVLURFQZO", "Medical,Food,Shopping,Subscription,Travel, Miscellaneous"],
    ["POS XXXXXXXXXXXX1111 DECATHLON SPORTS", "Medical,Food,Shopping,Subscription,Travel, Miscellaneous"],
    ["POS XXXXXXXXXXXX1111 IKEA INDIA PVT L", "Medical,Food,Shopping,Subscription,Travel, Miscellaneous"],
    ["POS XXXXXXXXXXXX1111 WWW AMAZON IN", "Medical,Food,Shopping,Subscription,Travel, Miscellaneous"],
    ["ME DC SI XXXXXXXXXXXX1111 SPOTIFY SI", "Medical,Food,Shopping,Subscription,Travel, Miscellaneous"],
    ["POS/NETFLIX/1140920002/100623/17:25", "Medical,Food,Shopping,Subscription,Travel, Miscellaneous"],
    ["POS XXXXXXXXXXXX1110 MAKEMYTRIP INDIA", "Medical,Food,Shopping,Subscription,Travel, Miscellaneous"],
    ["BIL/ONL/000691178015/IRCTC Serv/XZZBX91LTCY1AZ", "Medical,Food,Shopping,Subscription,Travel, Miscellaneous"]
]

#create interface
gui = gr.Interface(title=title,
                   description=description,
                   fn=zero_shot,
                   inputs=[preprocess(input1), input2],
                   outputs=[output],
                   examples=transactions_and_tags)

#display the interface
gui.launch()