bm25_test_v2 / app.py
dominiks's picture
next test
64632b0
import gradio as gr
import Stemmer # from PyStemmer
import bm25s
import numpy as np
import time
import os
examples = ["Can officers always order a passenger out of a car?","Find me briefs about credential searches", "Can police search an impounded car without a warrant?", "State is arguing State v. Carty is not good law"]
def show_user_query(user_message, history):
return "", history + [[user_message, None]]
def retrieve(history):
query = history[-1][0]
print ("query", query)
query_tokens = bm25s.tokenize(query, stemmer=stemmer)
results, scores = retriever.retrieve(query_tokens, k=10)
query_response = " ".join([str(i) for i in results])
#query_response = retrieve(history[-1][0])
#print (query_response)
history[-1][1] = ""
for character in query_response:
history[-1][1] += character
time.sleep(0.0001)
yield history
with gr.Blocks() as demo:
#gr.LoginButton()
chatbot = gr.Chatbot()
query_textbox = gr.Textbox()
examples = gr.Examples(examples, query_textbox)
next_button = gr.ClearButton(chatbot, value="Clear", visible=True)
query_textbox.submit(show_user_query, [query_textbox, chatbot], [query_textbox, chatbot], queue=False).then(retrieve, chatbot, chatbot)
def load_bm25():
stemmer = Stemmer.Stemmer("english")
retriever = bm25s.BM25.load("NJ_index_LLM_chunking", mmap=False)
return retriever, stemmer
retriever, stemmer = load_bm25()
#demo.launch(share=True, auth=(os.environ.get("username"), os.environ.get("password")))
#demo.launch(auth=("admin", "pass1234"))
demo.launch()