simiandb / app.py
fizban's picture
add placeholder
1e95982
raw
history blame contribute delete
No virus
949 Bytes
# -*- 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"))