import gradio as gr
import sentence_transformers
from sentence_transformers import SentenceTransformer
import torch
from sentence_transformers.util import semantic_search
import pandas as pd

model = SentenceTransformer('JoBeer/all-mpnet-base-v2-eclass')

corpus = pd.read_json('corpus.jsonl', lines = True, encoding = 'utf-8')

def predict(name, description):
    text = 'Description: '+ description + '; Name: ' + name 
    query_embedding = model.encode(text, convert_to_tensor=True)

    corpus_embeddings = torch.Tensor(corpus["embeddings"])

    output = sentence_transformers.util.semantic_search(query_embedding, corpus_embeddings, top_k = 5)

    preferedName1 = corpus.iloc[output[0][0].get('corpus_id'),2]
    definition1 = corpus.iloc[output[0][0].get('corpus_id'),1]
    IRDI1 = corpus.iloc[output[0][0].get('corpus_id'),4]
    score1 = output[0][0].get('score')

    preferedName2 = corpus.iloc[output[0][1].get('corpus_id'),2]
    definition2 = corpus.iloc[output[0][1].get('corpus_id'),1]
    IRDI2 = corpus.iloc[output[0][1].get('corpus_id'),4]
    score2 = output[0][1].get('score')

    df = [[preferedName1, definition1, IRDI1, score1], [preferedName2, definition2, IRDI1, score2]]
    
    return pd.Dataframe(df)

interface = gr.Interface(fn = predict, 
            inputs = [gr.Textbox(label="Name:", placeholder="z.B. GTIN", lines=1), gr.Textbox(label="Description:", placeholder="z.B. Globel Trade Item Number", lines=1)], 
            #outputs = [gr.Textbox(label = 'preferedName'),gr.Textbox(label = 'definition'), gr.Textbox(label = 'IDRI'),gr.Textbox(label = 'score')],
            outputs = [gr.Dataframe(row_count = (2, "fixed"), col_count=(4, "fixed"), label="Predictions", headers=['preferedName', 'definition', 'IRDI', 'score'])],
            examples = [['GTIN', 'Globel Trade Item Number'], ['Global Trade Item Number', 'the identification number from the GS1 system with which the trading units can be uniquely identified worldwide'],
                        ['Device type', 'describing a set of common specific characteristics in products or goods'], ['Item type','the type of product, an item can be assigned to'], 
                        ['Nominal power','power being consumed by or dissipated within an electric component as a variable'], ['Power consumption', 'power that is typically taken from the auxiliary power supply when the device is operating normally']], theme = 'huggingface',
            title = 'ECLASS-Property-Search')
    
interface.launch()