v1.1 only cbg search index

app.py

@@ -1,152 +1,152 @@
-import time
-import gradio as gr
-from datasets import load_dataset
-import pandas as pd
-from sentence_transformers import SentenceTransformer
-from sentence_transformers.quantization import quantize_embeddings
-import faiss
-from usearch.index import Index
-import datetime
-
-
-# Load titles and texts
-title_text_dataset = load_dataset("suanan/BP_POC", split="train", num_proc=4).select_columns(["url", "title", "text"])
-
-# Load the int8 and binary indices. Int8 is loaded as a view to save memory, as we never actually perform search with it.
-int8_view = Index.restore("index/BP_int8_usearch_1m.index", view=True)
-binary_index: faiss.IndexBinaryFlat = faiss.read_index_binary("index/BP_ubinary_faiss_1m.index")
-# binary_ivf: faiss.IndexBinaryIVF = faiss.read_index_binary("BP_ubinary_ivf_faiss_50m.index")
-
-# Load the SentenceTransformer model for embedding the queries
-model = SentenceTransformer(
-    "BAAI/bge-m3",
-    prompts={
-        "retrieval": "Represent this sentence for searching relevant passages: ",
-    },
-    default_prompt_name="retrieval",
-)
-
-
-def search(query, top_k: int = 100, rescore_multiplier: int = 1, use_approx: bool = False):
-    # 獲取當前時間
-    now = datetime.datetime.now()
-    print(f"當前時間: {now}, 問題: {query}")
-    # 1. Embed the query as float32
-    start_time = time.time()
-    query_embedding = model.encode(query)
-    embed_time = time.time() - start_time
-
-    # 2. Quantize the query to ubinary
-    start_time = time.time()
-    query_embedding_ubinary = quantize_embeddings(query_embedding.reshape(1, -1), "ubinary")
-    quantize_time = time.time() - start_time
-
-    # 3. Search the binary index (either exact or approximate)
-    # index = binary_ivf if use_approx else binary_index
-    index = binary_index
-    start_time = time.time()
-    _scores, binary_ids = index.search(query_embedding_ubinary, top_k * rescore_multiplier)
-    binary_ids = binary_ids[0]
-    search_time = time.time() - start_time
-
-    # 4. Load the corresponding int8 embeddings
-    start_time = time.time()
-    int8_embeddings = int8_view[binary_ids].astype(int)
-    load_time = time.time() - start_time
-
-    # 5. Rescore the top_k * rescore_multiplier using the float32 query embedding and the int8 document embeddings
-    start_time = time.time()
-    scores = query_embedding @ int8_embeddings.T
-    rescore_time = time.time() - start_time
-
-    # 6. Sort the scores and return the top_k
-    start_time = time.time()
-    indices = scores.argsort()[::-1][:top_k]
-    top_k_indices = binary_ids[indices]
-    top_k_scores = scores[indices]
-    top_k_urls, top_k_titles, top_k_texts = zip(
-        *[(title_text_dataset[idx]["url"], title_text_dataset[idx]["title"], title_text_dataset[idx]["text"]) for idx in top_k_indices.tolist()]
-    )
-    df = pd.DataFrame(
-        {"Score": [round(value, 2) for value in top_k_scores], "Url": top_k_urls, "Title": top_k_titles, "Text": top_k_texts}
-    )
-    sort_time = time.time() - start_time
-
-    return df, {
-        "Embed Time": f"{embed_time:.4f} s",
-        "Quantize Time": f"{quantize_time:.4f} s",
-        "Search Time": f"{search_time:.4f} s",
-        "Load Time": f"{load_time:.4f} s",
-        "Rescore Time": f"{rescore_time:.4f} s",
-        "Sort Time": f"{sort_time:.4f} s",
-        "Total search Time": f"{quantize_time + search_time + load_time + rescore_time + sort_time:.4f} s",
-    }
-def update_info(value):
-    return f"{value}筆顯示出來"
-
-with gr.Blocks(title="") as demo:
-    gr.Markdown(
-        """
-## 官網 Dataset & opensource model BAAI/bge-m3
-### v1 測試POC
-
-
-Details:
-1. 中文搜尋ok，英文像是:iphone 15，embedding的時候沒有轉成小寫，需要 寫成iPhone才可以準確搜尋到
-2. 環境資源: python 3.10, linux: ubuntu 22.04, only cpu, ram max:7.7GB min:4.5GB 使用以上資源
-3.
-
-建立步驟:
-1. excel 轉成 [dataset](https://huggingface.co/datasets/suanan/BP_POC), 花費約10秒內
-2. dataset 內 轉成 title & text 做 embedding，以後可以新增keyword來加強搜尋出來的結果排序往前
-3. 之後透過 Quantized Retrieval - Binary Search solution進行搜尋
-
-
-"""
-    )
-    with gr.Row():
-        with gr.Column(scale=75):
-            query = gr.Textbox(
-                label="官網 Dataset & opensource model BAAI/bge-m3,  v1 測試POC",
-                placeholder="輸入搜尋關鍵字或問句",
-            )
-        with gr.Column(scale=25):
-            use_approx = gr.Radio(
-                choices=[("精確搜尋", False), ("相關搜尋", True)],
-                value=False,
-                label="搜尋方法",
-            )
-
-    with gr.Row():
-        with gr.Column(scale=2):
-            top_k = gr.Slider(
-                minimum=10,
-                maximum=1000,
-                step=5,
-                value=100,
-                label="顯示搜尋前幾筆",
-            )
-            info_text = gr.Textbox(value=update_info(top_k.value), interactive=False)
-        with gr.Column(scale=2):
-            rescore_multiplier = gr.Slider(
-                minimum=1,
-                maximum=10,
-                step=1,
-                value=1,
-                label="Rescore multiplier",
-                info="Search for `rescore_multiplier` as many documents to rescore",
-            )
-
-    search_button = gr.Button(value="Search")
-
-    with gr.Row():
-        with gr.Column(scale=4):
-            output = gr.Dataframe(headers=["Score", "Title", "Text"])
-        with gr.Column(scale=1):
-            json = gr.JSON()
-    top_k.change(fn=update_info, inputs=top_k, outputs=info_text)
-    query.submit(search, inputs=[query, top_k, rescore_multiplier, use_approx], outputs=[output, json])
-    search_button.click(search, inputs=[query, top_k, rescore_multiplier, use_approx], outputs=[output, json])
-
-demo.queue()
-demo.launch(share=True)
+import time
+import gradio as gr
+from datasets import load_dataset
+import pandas as pd
+from sentence_transformers import SentenceTransformer
+from sentence_transformers.quantization import quantize_embeddings
+import faiss
+from usearch.index import Index
+import datetime
+
+
+# Load titles and texts
+title_text_dataset = load_dataset("suanan/BP_POC", split="train", num_proc=4).select_columns(["url", "title", "text"])
+
+# Load the int8 and binary indices. Int8 is loaded as a view to save memory, as we never actually perform search with it.
+int8_view = Index.restore("index/BP_CBG_int8_usearch_1m.index", view=True)
+binary_index: faiss.IndexBinaryFlat = faiss.read_index_binary("index/BP_CBG_ubinary_faiss_1m.index")
+# binary_ivf: faiss.IndexBinaryIVF = faiss.read_index_binary("BP_ubinary_ivf_faiss_50m.index")
+
+# Load the SentenceTransformer model for embedding the queries
+model = SentenceTransformer(
+    "BAAI/bge-m3",
+    prompts={
+        "retrieval": "Represent this sentence for searching relevant passages: ",
+    },
+    default_prompt_name="retrieval",
+)
+
+
+def search(query, top_k: int = 100, rescore_multiplier: int = 1, use_approx: bool = False):
+    # 獲取當前時間
+    now = datetime.datetime.now()
+    print(f"當前時間: {now}, 問題: {query}")
+    # 1. Embed the query as float32
+    start_time = time.time()
+    query_embedding = model.encode(query)
+    embed_time = time.time() - start_time
+
+    # 2. Quantize the query to ubinary
+    start_time = time.time()
+    query_embedding_ubinary = quantize_embeddings(query_embedding.reshape(1, -1), "ubinary")
+    quantize_time = time.time() - start_time
+
+    # 3. Search the binary index (either exact or approximate)
+    # index = binary_ivf if use_approx else binary_index
+    index = binary_index
+    start_time = time.time()
+    _scores, binary_ids = index.search(query_embedding_ubinary, top_k * rescore_multiplier)
+    binary_ids = binary_ids[0]
+    search_time = time.time() - start_time
+
+    # 4. Load the corresponding int8 embeddings
+    start_time = time.time()
+    int8_embeddings = int8_view[binary_ids].astype(int)
+    load_time = time.time() - start_time
+
+    # 5. Rescore the top_k * rescore_multiplier using the float32 query embedding and the int8 document embeddings
+    start_time = time.time()
+    scores = query_embedding @ int8_embeddings.T
+    rescore_time = time.time() - start_time
+
+    # 6. Sort the scores and return the top_k
+    start_time = time.time()
+    indices = scores.argsort()[::-1][:top_k]
+    top_k_indices = binary_ids[indices]
+    top_k_scores = scores[indices]
+    top_k_urls, top_k_titles, top_k_texts = zip(
+        *[(title_text_dataset[idx]["url"], title_text_dataset[idx]["title"], title_text_dataset[idx]["text"]) for idx in top_k_indices.tolist()]
+    )
+    df = pd.DataFrame(
+        {"Score": [round(value, 2) for value in top_k_scores], "Url": top_k_urls, "Title": top_k_titles, "Text": top_k_texts}
+    )
+    sort_time = time.time() - start_time
+
+    return df, {
+        "Embed Time": f"{embed_time:.4f} s",
+        "Quantize Time": f"{quantize_time:.4f} s",
+        "Search Time": f"{search_time:.4f} s",
+        "Load Time": f"{load_time:.4f} s",
+        "Rescore Time": f"{rescore_time:.4f} s",
+        "Sort Time": f"{sort_time:.4f} s",
+        "Total search Time": f"{quantize_time + search_time + load_time + rescore_time + sort_time:.4f} s",
+    }
+def update_info(value):
+    return f"{value}筆顯示出來"
+
+with gr.Blocks(title="") as demo:
+    gr.Markdown(
+        """
+## 官網 Dataset & opensource model BAAI/bge-m3
+### v1 測試POC
+
+
+Details:
+1. 中文搜尋ok，英文像是:iphone 15，embedding的時候沒有轉成小寫，需要 寫成iPhone才可以準確搜尋到
+2. 環境資源: python 3.10, linux: ubuntu 22.04, only cpu, ram max:7.7GB min:4.5GB 使用以上資源
+3.
+
+建立步驟:
+1. excel 轉成 [dataset](https://huggingface.co/datasets/suanan/BP_POC), 花費約10秒內
+2. dataset 內 轉成 title & text 做 embedding，以後可以新增keyword來加強搜��出來的結果排序往前
+3. 之後透過 Quantized Retrieval - Binary Search solution進行搜尋
+
+
+"""
+    )
+    with gr.Row():
+        with gr.Column(scale=75):
+            query = gr.Textbox(
+                label="官網 Dataset & opensource model BAAI/bge-m3,  v1 測試POC",
+                placeholder="輸入搜尋關鍵字或問句",
+            )
+        with gr.Column(scale=25):
+            use_approx = gr.Radio(
+                choices=[("精確搜尋", False), ("相關搜尋", True)],
+                value=False,
+                label="搜尋方法",
+            )
+
+    with gr.Row():
+        with gr.Column(scale=2):
+            top_k = gr.Slider(
+                minimum=10,
+                maximum=1000,
+                step=5,
+                value=100,
+                label="顯示搜尋前幾筆",
+            )
+            info_text = gr.Textbox(value=update_info(top_k.value), interactive=False)
+        with gr.Column(scale=2):
+            rescore_multiplier = gr.Slider(
+                minimum=1,
+                maximum=10,
+                step=1,
+                value=1,
+                label="Rescore multiplier",
+                info="Search for `rescore_multiplier` as many documents to rescore",
+            )
+
+    search_button = gr.Button(value="Search")
+
+    with gr.Row():
+        with gr.Column(scale=4):
+            output = gr.Dataframe(headers=["Score", "Title", "Text"])
+        with gr.Column(scale=1):
+            json = gr.JSON()
+    top_k.change(fn=update_info, inputs=top_k, outputs=info_text)
+    query.submit(search, inputs=[query, top_k, rescore_multiplier, use_approx], outputs=[output, json])
+    search_button.click(search, inputs=[query, top_k, rescore_multiplier, use_approx], outputs=[output, json])
+
+demo.queue()
+demo.launch(share=True)