import numpy as np
import gradio as gr
from sentence_transformers import SentenceTransformer


minilm = SentenceTransformer('all-MiniLM-L12-v2')
roberta = SentenceTransformer('all-distilroberta-v1')
glove = SentenceTransformer('average_word_embeddings_glove.840B.300d')


labels = ["contradiction", "entailment", "neutral"]

def predict(sentence1, sentence2):
    sentence_pairs = np.array([[str(sentence1), str(sentence2)]])
    print(sentence1)
    print(sentence2)
#    test_data = BertSemanticDataGenerator(
#        sentence_pairs, labels=None, batch_size=1, shuffle=False, include_targets=False,
#    )
#    probs = model.predict(test_data[0])[0]
    
#    labels_probs = {labels[i]: float(probs[i]) for i, _ in enumerate(labels)}
 #   return labels_probs

examples = [["Two women are observing something together.", "Two women are standing with their eyes closed."],
            ["A smiling costumed woman is holding an umbrella", "A happy woman in a fairy costume holds an umbrella"],
            ["A soccer game with multiple males playing", "Some men are playing a sport"],            
            ]   

gr.Interface(
    fn=predict,
    title="Semantic Song Search",
    description = "Search for songs based on the meaning in the song's lyrics using a variety of embeddings",
    inputs=["text", "text"],
    examples=examples,
    #outputs=gr.Textbox(label='Prediction'),
    outputs=gr.outputs.Label(num_top_classes=3, label='Semantic similarity'),
    cache_examples=True,
    article = "Author: @sheacon",
).launch(debug=True, enable_queue=True)