Upload folder using huggingface_hub

app.py

@@ -7,6 +7,7 @@ import gradio as gr
 import faiss
 import numpy as np
 import torch
+from tqdm import tqdm
 from sentence_transformers import SentenceTransformer
 
 
@@ -17,7 +18,9 @@ file_example = """Please upload a JSON file with a "text" field (with optional "
     {"title": "Title A", "text": "This an example text with the title"},
     {"title": "Title B", "text": "This an example text with the title"},
 ]
-```"""
+```
+Due to the computation resources, please test with small scale data.
+"""
 
 
 def create_index(embeddings, use_gpu):
@@ -42,13 +45,26 @@ def upload_file_fn(
         documents = []
         for obj in document_data:
             text = obj["title"] + "\n" + obj["text"] if obj.get("title") else obj["text"]
-            documents.append(text)
+            if len(str(text).strip()):
+                documents.append(text)
+            else:
+                documents.append(model.tokenizer.eos_token)
     except Exception as e:
         print(e)
-        gr.Warning("Read the file failed. Please check the data format.")
-        return None, None
+        gr.Error("Read the file failed. Please check the data format.")
+        gr.Error(str(e))
+        return None
+
+    if len(documents) < 3:
+        gr.Error("Please upload more than 3 documents.")
+        return None
+
+    gr.Info(f"Upload {len(documents)} documents.")
+    if len(documents) > 2000:
+        gr.Info(f"Cut uploaded documents to 2000.")
+        documents = documents[: 2000]
     
-    documents_embeddings = model.encode(documents)
+    documents_embeddings = model.encode(documents, show_progress_bar=True)
 
     document_index = create_index(documents_embeddings, use_gpu=False)
     
@@ -56,25 +72,30 @@ def upload_file_fn(
         torch.cuda.empty_cache()
         torch.cuda.ipc_collect()
     
-    return document_index, document_data
+    print("upload is OK")
+    document_state = {"document_data": document_data, "document_index": document_index}
+    return document_state,
 
 
 def clear_file_fn():
-    return None, None
+    return None
 
 
-def retrieve_document_fn(question, instruct, document_states):
-    document_data, document_index = document_states
+def retrieve_document_fn(question, document_states, instruct):
     num_retrieval_doc = 3
-    if document_index is None or document_data is None:
+
+    if document_states is None:
         gr.Warning("Please upload documents first!")
-        return [None for i in range(num_retrieval_doc)]
+        return [None for i in range(num_retrieval_doc)] + [None]
+    
+    print(document_states)
+    document_data, document_index = document_states["document_data"], document_states["document_index"]
     
-    question_embedding = model.encode([instruct + question])
+    question_embedding = model.encode([str(instruct) + str(question)])
     batch_scores, batch_inxs = document_index.search(question_embedding, k=num_retrieval_doc)
     
     answers = [document_data[i]["text"] for i in batch_inxs[0][:num_retrieval_doc]]
-    return tuple(answers)
+    return answers[0], answers[1], answers[2], document_states
 
 
 def main(args):
@@ -82,10 +103,8 @@ def main(args):
 
     model = SentenceTransformer(args.model_name_or_path)
 
-    document_index = gr.State()
-    document_data = gr.State()
+    document_state = gr.State()
 
-    
     with open(Path(__file__).parent / "resources/head.html") as html_file:
         head = html_file.read().strip()
     with gr.Blocks(theme=gr.themes.Soft(font="sans-serif").set(background_fill_primary="linear-gradient(90deg, #e3ffe7 0%, #d9e7ff 100%)", background_fill_primary_dark="linear-gradient(90deg, #4b6cb7 0%, #182848 100%)",),
@@ -98,23 +117,24 @@ def main(args):
         doc_files_box = gr.File(label="Upload Documents", file_types=[".json"], file_count="single")
         retrieval_interface = gr.Interface(
             fn=retrieve_document_fn,
-            inputs=["text"],
-            outputs=["text", "text",  "text"],
-            additional_inputs=[gr.Textbox("Instruct: Given a query, retrieve documents that answer the query. \n Query: ", label="Instruct of Query"), gr.State(value=[document_data, document_index])],
+            inputs=["text", document_state],
+            outputs=[gr.Text(label="Recall-1"), gr.Text(label="Recall-2"),  gr.Text(label="Recall-3"), gr.State()],
+            additional_inputs=[gr.Textbox("Instruct: Given a query, retrieve documents that answer the query. \n Query: ", label="Instruct of Query")],
             concurrency_limit=1,
         )
+        
 
         doc_files_box.upload(
             upload_file_fn, 
             [doc_files_box], 
-            [document_index, document_data],
+            [document_state],
             queue=True,
             trigger_mode="once"
         )
         doc_files_box.clear(
-                upload_file_fn, 
+                clear_file_fn, 
                 None, 
-                [document_index, document_data],
+                [document_state],
                 queue=True,
                 trigger_mode="once"
             )
@@ -123,7 +143,8 @@ def main(args):
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser.add_argument("--model_name_or_path", type=str, default="HIT-TMG/KaLM-embedding-multilingual-mini-instruct-v1.5")
+    # parser.add_argument("--model_name_or_path", type=str, default="HIT-TMG/KaLM-embedding-multilingual-mini-instruct-v1.5")
+    parser.add_argument("--model_name_or_path", type=str, default="/raid/hxs/Checkpoints/huggingface_models/bge-base-en-v1.5")
 
     args = parser.parse_args()
     main(args)