| |
| """ |
| 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() |
|
|
| |
| |
