import gradio as gr
import numpy as np
import pandas as pd
from sentence_transformers import SentenceTransformer, util

model = SentenceTransformer('all-MiniLM-L6-v2')

sentences = ["Address", "Port", "Country", "Name", "Quantity"]

embeddings = model.encode(sentences)
def greet(key):
    emb = model.encode([key])
    cosine_scores = np.array(util.cos_sim(embeddings, emb))
    print(np.array(cosine_scores))
    print(np.argmax(cosine_scores))
    # Find the pairs with the highest cosine similarity scores
    df = pd.DataFrame(data=list(zip(sentences, cosine_scores[:, 0])), columns=["Entity", "Score"])
    df = df.sort_values(by=['Score'], ascending=False)
    return df


demo = gr.Interface(fn=greet, inputs="text", outputs="dataframe", title="Key Matcher",
                    description=f"Find the best match for a key from the list of {sentences}",
                    examples=[["Addres"], ["Ship To"], ["Shipping Name"], ["Amount"]],
                    live=True)

demo.launch()