Shakshi3104 commited on
Commit
8a06424
β€’
1 Parent(s): 481da03

[add] Implement UI with gradio

Browse files
Files changed (1) hide show
  1. app.py +50 -0
app.py ADDED
@@ -0,0 +1,50 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import pandas as pd
3
+
4
+ from model.search.hybrid import HybridSearchClient
5
+ from model.data.notion_db import fetch_sakurap_corpus
6
+
7
+
8
+ def search(search_client: HybridSearchClient):
9
+ def _search(query: str) -> pd.DataFrame:
10
+ results = search_client.search_top_n(query)
11
+ result = results[0]
12
+ result["rank"] = result["rank"] + 1
13
+ result = result[["rank", "title", "content", "rank_sparse", "rank_dense"]]
14
+ result.columns = ["rank", "title", "rap lyric", "rank: surface", "rank: vector"]
15
+ return result
16
+
17
+ return _search
18
+
19
+
20
+ if __name__ == "__main__":
21
+ # Load dataset
22
+ sakurap_df = fetch_sakurap_corpus("./data/sakurap_corpus.csv")
23
+ # Initialize search client
24
+ search_client = HybridSearchClient.from_dataframe(sakurap_df, "content")
25
+
26
+ with gr.Blocks() as search_interface:
27
+ gr.Markdown("""
28
+ # πŸ’Ž Cobalt
29
+ Demo app for hybrid search with vector and surface search using [Ruri](https://huggingface.co/cl-nagoya/ruri-large), [BM25](https://github.com/dorianbrown/rank_bm25) and [Voyager](https://spotify.github.io/voyager/).
30
+ """)
31
+ # Input query
32
+ search_query = gr.Textbox(label="Query", submit_btn=True)
33
+
34
+ gr.Markdown("""
35
+ ## Search Results
36
+
37
+ """)
38
+ # Search result
39
+ result_table = gr.DataFrame(label="Result",
40
+ column_widths=["5%", "20%", "65%", "5%", "5%"],
41
+ wrap=True,
42
+ datatype=["str", "str", "markdown", "str", "str"],
43
+ interactive=False)
44
+
45
+ # Event handler
46
+ search_query.submit(fn=search(search_client), inputs=search_query, outputs=result_table)
47
+
48
+ # App launch
49
+ search_interface.queue()
50
+ search_interface.launch()