Update app.py

app.py

@@ -4,15 +4,22 @@ import torch.nn.functional as F
 import hnswlib
 import gradio as gr
 import numpy as np
+import json
+import datetime
 
 seperator = "-HFSEP-"
 base_name="intfloat/e5-large-v2"
 device="cuda"
 max_length=512
+max_batch_size = 500
 tokenizer = AutoTokenizer.from_pretrained(base_name)
 model = AutoModel.from_pretrained(base_name).to(device)
 
+def current_timestamp():
+    return datetime.datetime.utcnow().timestamp()
+
 def get_embeddings(input_texts):
+    input_texts = input_texts[:max_batch_size]
     batch_dict = tokenizer(
         input_texts,
         max_length=max_length,
@@ -52,16 +59,18 @@ def create_hnsw_index(embeddings_np, space='ip', ef_construction=100, M=16):
 def gradio_function(query, paragraph_chunks, top_k):
     paragraph_chunks = paragraph_chunks.split(seperator)  # Split the comma-separated values into a list
     paragraph_chunks = [item.strip() for item in paragraph_chunks]  # Trim whitespace from each item
+    print("Len of batches", len(paragraph_chunks))
 
-    print("creating embeddings")
+    print("creating embeddings", current_timestamp())
     embeddings_np = get_embeddings([query]+paragraph_chunks)
     query_embedding, chunks_embeddings = embeddings_np[0], embeddings_np[1:]
     
-    print("creating index")
+    print("creating index", current_timestamp())
     search_index = create_hnsw_index(chunks_embeddings)
-    print("searching index")
+    print("searching index", current_timestamp())
     labels, _ = search_index.knn_query(query_embedding, k=min(int(top_k), len(chunks_embeddings)))
-    return f"The closes labels are: {labels}"
+    labels = labels[0].tolist()
+    return json.dumps(labels)
 
 interface = gr.Interface(
     fn=gradio_function, 
@@ -73,4 +82,4 @@ interface = gr.Interface(
     outputs="text"
 )
 
-interface.launch()
+interface.launch()