Upload folder using huggingface_hub

app.py

@@ -53,18 +53,24 @@ def upload_file_fn(
         print(e)
         gr.Error("Read the file failed. Please check the data format.")
         gr.Error(str(e))
-        return None
+        return None, gr.update(interactive=False)
 
     if len(documents) < 3:
         gr.Error("Please upload more than 3 documents.")
-        return None
+        return None, gr.update(interactive=False)
 
     gr.Info(f"Upload {len(documents)} documents.")
-    if len(documents) > 2000:
-        gr.Info(f"Cut uploaded documents to 2000.")
-        documents = documents[: 2000]
+    if len(documents) > 1000:
+        gr.Info(f"Cut uploaded documents to 1000 due to the computation resource.")
+        documents = documents[: 1000]
     
-    documents_embeddings = model.encode(documents, show_progress_bar=True)
+    # documents_embeddings = model.encode(documents, show_progress_bar=True)
+    documents_embeddings = []
+    batch_size = 8
+    for i in tqdm(range(0, len(documents), batch_size)):
+        batch_documents = documents[i: i+batch_size]
+        batch_embeddings = model.encode(batch_documents, show_progress_bar=True)
+        documents_embeddings.extend(batch_embeddings)
 
     document_index = create_index(documents_embeddings, use_gpu=False)
     
@@ -72,13 +78,12 @@ def upload_file_fn(
         torch.cuda.empty_cache()
         torch.cuda.ipc_collect()
     
-    print("upload is OK")
     document_state = {"document_data": document_data, "document_index": document_index}
-    return document_state,
+    return document_state, gr.update(interactive=True)
 
 
 def clear_file_fn():
-    return None
+    return None, gr.update(interactive=True)
 
 
 def retrieve_document_fn(question, document_states, instruct):
@@ -87,12 +92,11 @@ def retrieve_document_fn(question, document_states, instruct):
     if document_states is None:
         gr.Warning("Please upload documents first!")
         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([str(instruct) + str(question)])
-    batch_scores, batch_inxs = document_index.search(question_embedding, k=num_retrieval_doc)
+    batch_scores, batch_inxs = document_index.search(question_embedding, k=min(len(document_data), 150))
     
     answers = [document_data[i]["text"] for i in batch_inxs[0][:num_retrieval_doc]]
     return answers[0], answers[1], answers[2], document_states
@@ -101,7 +105,10 @@ def retrieve_document_fn(question, document_states, instruct):
 def main(args):
     global model
 
-    model = SentenceTransformer(args.model_name_or_path)
+    model = SentenceTransformer(
+        args.model_name_or_path, 
+        revision=args.revision,
+        backend="openvino")
 
     document_state = gr.State()
 
@@ -117,24 +124,25 @@ 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", document_state],
+            inputs=[gr.Textbox(label="Query"), 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")],
+            additional_inputs=[gr.Textbox("Instruct: Given a query, retrieve documents that answer the query. \n Query: ", label="Instruct of Query", lines=2)],
             concurrency_limit=1,
         )
+        # retrieval_interface.input_components[0] = gr.update(interactive=False)
         
 
         doc_files_box.upload(
             upload_file_fn, 
             [doc_files_box], 
-            [document_state],
+            [document_state, retrieval_interface.input_components[0]],
             queue=True,
             trigger_mode="once"
         )
         doc_files_box.clear(
                 clear_file_fn, 
                 None, 
-                [document_state],
+                [document_state, retrieval_interface.input_components[0]],
                 queue=True,
                 trigger_mode="once"
             )
@@ -145,6 +153,7 @@ 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="/raid/hxs/Checkpoints/huggingface_models/bge-base-en-v1.5")
+    parser.add_argument("--revision", type=str, default="refs/pr/2")
 
     args = parser.parse_args()
     main(args)