import numpy as np # linear algebra
import os
from sklearn.metrics.pairwise import cosine_similarity
from sentence_transformers import SentenceTransformer
import pickle
import gradio as gr

with open('embeddings.pkl', "rb") as fIn:
    stored_data = pickle.load(fIn)
    embeddings = stored_data['embeddings']

with open('answers.pkl', "rb") as fIn:
    stored_data = pickle.load(fIn)
    answers = stored_data['answers']

model = SentenceTransformer('msmarco-MiniLM-L-6-v3')

def predict(ques):
    embeded_ques = model.encode(ques)
    em_vec = np.vstack([embeded_ques]*10)
    scores = cosine_similarity(embeddings, em_vec)
    idx = np.argmax([np.mean(arr) for arr in scores])
    return answers[idx]

examples = ['What is the date of his death?', 'Did Einstein have siblings?', 'Who was his wife?', 'What was Einstein\'s father\'s name?', 'At what institutions did he study?']
inputs = gr.Textbox(label='query')
outputs = gr.Textbox(label='Answers')
title = "Similar faq"
description = "Retreive answers of similar queries using sentence transformers"
gr.Interface(fn = predict, inputs = inputs, examples=examples, outputs = outputs, title = title, description = description).launch()