dockerify the app

Dockerfile

@@ -0,0 +1,11 @@
+FROM python:3.9
+
+WORKDIR /code
+
+COPY ./requirements.txt /code/requirements.txt
+
+RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
+
+COPY . .
+
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

README.md

@@ -3,9 +3,7 @@ title: Embeddings Similarity
 emoji: 📚
 colorFrom: purple
 colorTo: gray
-sdk: gradio
-sdk_version: 3.41.2
-app_file: app.py
+sdk: docker
 pinned: false
 ---
 

app.py

@@ -2,10 +2,16 @@ from transformers import AutoTokenizer, AutoModel
 import torch
 import torch.nn.functional as F
 import hnswlib
-import gradio as gr
 import numpy as np
-import json
 import datetime
+from fastapi import FastAPI
+from pydantic import BaseModel
+from typing import List
+
+if torch.cuda.is_available():
+    print("CUDA is available! Inference on GPU!")
+else:
+    print("CUDA is not available. Inference on CPU.")
 
 seperator = "-HFSEP-"
 base_name="intfloat/e5-large-v2"
@@ -56,30 +62,24 @@ def create_hnsw_index(embeddings_np, space='ip', ef_construction=100, M=16):
     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("Len of batches", len(paragraph_chunks))
+app = FastAPI()
+
+class EmbeddingsSimilarityReq(BaseModel):
+    paragraphs: List[str]
+    query: str
+    top_k: int
+
+@app.post("/")
+async def find_similar_paragraphsitem(req: EmbeddingsSimilarityReq):
+    print("Len of batches", len(req.paragraphs))
 
     print("creating embeddings", current_timestamp())
-    embeddings_np = get_embeddings([query]+paragraph_chunks)
+    embeddings_np = get_embeddings([req.query]+req.paragraphs)
     query_embedding, chunks_embeddings = embeddings_np[0], embeddings_np[1:]
     
     print("creating index", current_timestamp())
     search_index = create_hnsw_index(chunks_embeddings)
     print("searching index", current_timestamp())
-    labels, _ = search_index.knn_query(query_embedding, k=min(int(top_k), len(chunks_embeddings)))
+    labels, _ = search_index.knn_query(query_embedding, k=min(int(req.top_k), len(chunks_embeddings)))
     labels = labels[0].tolist()
-    return json.dumps(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()
+    return labels

requirements.txt

@@ -1,4 +1,6 @@
 torch==2.0.1
 transformers
 gradio
-hnswlib
+hnswlib
+fastapi
+uvicorn[standard]