Update etown_mxbai/router.py

etown_mxbai/router.py

@@ -31,13 +31,30 @@ app.add_middleware(
 def inference(item: Validation):
     try:
         start_time = time.time()
-        embeddings = list(sparse_model.embed(item.prompt, batch_size=32))   # Assuming 'model' is defined elsewhere
+        embeddings = list(sparse_model.embed(item.prompt, batch_size=5))   # Assuming 'model' is defined elsewhere
+
+        serializable_embeddings = []
+        for embedding in embeddings:
+            # Assuming embedding object has attributes values and indices
+            if isinstance(embedding, SparseEmbedding):
+                values = embedding.values
+                indices = embedding.indices
+                serializable_embeddings.append({
+                    "values": values.tolist() if isinstance(values, np.ndarray) else values,
+                    "indices": indices.tolist() if isinstance(indices, np.ndarray) else indices
+                })
+            else:
+                # Fallback for other types of embeddings
+                serializable_embeddings.append({
+                    "values": embedding.tolist() if isinstance(embedding, np.ndarray) else str(embedding),
+                    "indices": list(range(len(embedding))) if isinstance(embedding, (np.ndarray, list)) else []
+                })
                 
         end_time = time.time()
         time_taken = end_time - start_time  # Calculate the time taken
         
         return JSONResponse(content={
-                "embeddings": embeddings,
+                "embeddings": serializable_embeddings,
                 "time_taken": f"{time_taken:.2f} seconds",
                 "Number_of_sentence_processed": len(item.prompt),  # Assuming you want to count words, not characters
                 "Model_response_space" : "prithvida/Splade_PP_en_v1",