# -*- coding: utf-8 -*-
"""
Created on Wed Mar 22 19:59:54 2023

"""

import gradio as gr
from simiandb import Simiandb
from langchain.embeddings import HuggingFaceEmbeddings
from sentence_transformers import CrossEncoder




model_name = "all-MiniLM-L6-v2"
hf = HuggingFaceEmbeddings(model_name=model_name)
cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')

documentdb = Simiandb("mystore", embedding_function=hf, mode="a")

def search(query):
    hits = documentdb.similarity_search(query, k=10)
    cross_inp = [[query, hit] for hit in hits]
    cross_scores = cross_encoder.predict(cross_inp)
    hits = [hit for _, hit in sorted(zip(cross_scores, hits), reverse=True)]
    return hits[0]

iface = gr.Interface(fn=search, inputs=gr.Textbox(lines=2, placeholder="Write a question to the Wikipedia..."), outputs="text")
iface.launch()

#print(search("what is the balloon boy hoax"))
# print(search("date of birth of elon musk"))