import gradio as gr
import spacy
import classy_classification
import csv
import pandas as pd
import os

title = "Few-Shot Text Classification with spaCy"
description = "classy-classification"

#define a function to process your input and output
def few_shot(doc, csv_file):
    df = pd.read_csv(csv_file.name)
    data = {}
    sample_size = 10

    candidate_labels = df['label'].unique().tolist()

    for label in candidate_labels:
        candidate_values = df.query(f"`label` == '{label}'").sample(
            n=sample_size)['text'].values.tolist()
        data[label] = candidate_values
   
    nlp = spacy.blank("en")
    nlp.add_pipe(
        "text_categorizer",
        config={
            "data": data,
            "model": "sentence-transformers/all-mpnet-base-v2",
            "device": "gpu"
        }
    )

    dictionary = nlp(doc)._.cats
    return dictionary

#create input and output objects
#input object1
input1 = gr.Textbox(label="Text")
#input object2
input2 = gr.File(label="CSV File")
#output object
output = gr.Label(label="Output")
#example object
examples = [
            ["Fertiglobe plc produces and distributes nitrogen fertilizers. The company offers ammonia used as a building block for other fertilizer products to industrial and agricultural customers; urea for agricultural and industrial customers; and diesel exhaust fluids for industrial customers. It operates in Europe, North America, South America, the Middle East, Africa, Asia, and Oceania. The company was incorporated in 2019 and is headquartered in Abu Dhabi, the United Arab Emirates.", os.path.join(os.path.dirname(__file__),"files/train1.csv")],
            ["As of December 18, 2009, Media of Medias Public Co. Ltd. was acquired by Grand Canal Land Public Company Limited, in a reverse merger transaction. Media of Medias Public Company Limited, together with its subsidiaries, engages in producing television programs and selling advertising air time, and property development businesses. It also involves in entertainment and recreation business, which include television, VCD business, selling pocket books, and operation of golf course and hotel. The company also engages in the development, construction, rental, and sale of condominiums and plazas. Media of Medias Public Company Limited is based in Bangkok, Thailand.", os.path.join(os.path.dirname(__file__),"files/train2.csv")],
            ["Milsy A.s. provides dairy products. It offers fresh dairy products, steamed unsmoked and smoked cheeses, and spreads of various flavors. The company was founded in 1954 and is based in Banovce nad Bebravou, Slovakia.", os.path.join(os.path.dirname(__file__),"files/train3.csv")]
            ]
#create interface
gui = gr.Interface(title=title,
                   description=description,
                   fn=few_shot,
                   inputs=[input1, input2],
                   outputs=[output],
                   examples=examples)

#display the interface
gui.launch()