Create app.py

app.py

@@ -0,0 +1,76 @@
+from transformers import AutoTokenizer, AutoModel
+import torch
+import torch.nn.functional as F
+import hnswlib
+import gradio as gr
+import numpy as np
+
+seperator = "-HFSEP-"
+base_name="intfloat/e5-large-v2"
+device="cuda"
+max_length=512
+tokenizer = AutoTokenizer.from_pretrained(base_name)
+model = AutoModel.from_pretrained(base_name).to(device)
+
+def get_embeddings(input_texts):
+    batch_dict = tokenizer(
+        input_texts,
+        max_length=max_length,
+        padding=True,
+        truncation=True,
+        return_tensors='pt'
+    ).to(device)
+
+    with torch.no_grad():
+        outputs = model(**batch_dict)
+        embeddings = _average_pool(
+            outputs.last_hidden_state, batch_dict['attention_mask']
+        )
+        embeddings = F.normalize(embeddings, p=2, dim=1)
+        embeddings_np = embeddings.cpu().numpy()
+        
+        if device == "cuda":
+            del embeddings
+            torch.cuda.empty_cache()        
+        
+        return embeddings_np
+
+def _average_pool(
+    last_hidden_states,
+    attention_mask
+):
+    last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
+    return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
+
+def create_hnsw_index(embeddings_np, space='ip', ef_construction=100, M=16):
+    index = hnswlib.Index(space=space, dim=len(embeddings_np[0]))
+    index.init_index(max_elements=len(embeddings_np), ef_construction=ef_construction, M=M)
+    ids = np.arange(embeddings_np.shape[0])
+    index.add_items(embeddings_np, ids)
+    return index
+
+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("creating embeddings")
+    embeddings_np = get_embeddings([query]+paragraph_chunks)
+    query_embedding, chunks_embeddings = embeddings_np[0], embeddings_np[1:]
+    
+    print("creating index")
+    search_index = create_hnsw_index(chunks_embeddings)
+    print("searching index")
+    labels, _ = search_index.knn_query(query_embedding, k=min(int(top_k), len(chunks_embeddings)))
+    return f"The closes labels are: {labels}"
+
+interface = gr.Interface(
+    fn=gradio_function, 
+    inputs=[
+        gr.Textbox(placeholder="Enter a user query..."),
+        gr.Textbox(placeholder="Enter comma-separated strings..."),
+        gr.Number()
+    ],
+    outputs="text"
+)
+
+interface.launch()