|
|
|
""" |
|
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) |
|
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="text", outputs="text") |
|
iface.launch() |
|
|
|
|
