merge

app.py

@@ -1,26 +1,94 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
 from sentence_transformers import SentenceTransformer
+from transformers import AutoTokenizer, AutoModelForMaskedLM
 from typing import List
+import torch
+from functools import lru_cache
+import logging
 
-# Load the pre-trained sentence transformer model
+# 🔧 Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# 🚀 Initialize FastAPI app
+app = FastAPI()
+logger.info("Starting FastAPI application")
+
+# 🔌 Load SentenceTransformer models
+logger.info("Loading BGE small model...")
 bge_small_model = SentenceTransformer('BAAI/bge-small-en-v1.5', device="cpu")
+logger.info("Loaded BGE small model")
+
+logger.info("Loading All-MPNet model...")
 all_mp_net_model = SentenceTransformer('sentence-transformers/all-mpnet-base-v2', device="cpu")
+logger.info("Loaded All-MPNet model")
 
-# Initialize FastAPI app
-app = FastAPI()
+# 🔌 Load SPLADE model
+logger.info("Loading SPLADE model...")
+SPLADE_MODEL = AutoModelForMaskedLM.from_pretrained("naver/splade-cocondenser-ensembledistil", trust_remote_code=True)
+SPLADE_TOKENIZER = AutoTokenizer.from_pretrained("naver/splade-cocondenser-ensembledistil")
+SPLADE_MODEL.eval()
+logger.info("Loaded SPLADE model")
 
-# Request body model
+# 📦 Request and response models
 class TextInput(BaseModel):
-    text: List[str]  # List of sentences or text data
+    text: List[str]
     model_name: str
 
-# Route to calculate embeddings
+class SparseVector(BaseModel):
+    indices: List[int]
+    values: List[float]
+
+# 🧠 LRU cacheable versions
+@lru_cache(maxsize=1000)
+def encode_dense_cached(model_name: str, text: str):
+    logger.info(f"Encoding dense text with model {model_name}: {text}")
+    if model_name == "BM":
+        embedding = all_mp_net_model.encode([text])[0].tolist()
+    else:
+        embedding = bge_small_model.encode([text])[0].tolist()
+    logger.info(f"Finished encoding dense text")
+    return embedding
+
+@lru_cache(maxsize=1000)
+def encode_splade_cached(text: str) -> SparseVector:
+    logger.info(f"Encoding SPLADE sparse vector: {text}")
+    inputs = SPLADE_TOKENIZER(text, return_tensors="pt", truncation=True)
+    with torch.no_grad():
+        outputs = SPLADE_MODEL(**inputs)
+
+    logits = outputs.logits[0]
+    relu_log = torch.log1p(torch.relu(logits))
+    nonzero = relu_log.nonzero(as_tuple=False)
+
+    if nonzero.shape[0] == 0:
+        logger.info("No non-zero values found in SPLADE output")
+        return SparseVector(indices=[], values=[])
+
+    vocab_indices = nonzero[:, 1]
+    values = relu_log[nonzero[:, 0], nonzero[:, 1]]
+
+    logger.info(f"SPLADE encoding complete with {len(vocab_indices)} dimensions")
+    return SparseVector(
+        indices=vocab_indices.cpu().numpy().tolist(),
+        values=values.cpu().numpy().tolist()
+    )
+
+# 🚀 Main endpoint
 @app.post("/get-embedding/")
 async def get_embedding(input: TextInput):
-    # Generate embeddings using the sentence transformer model
-    if input.model_name == "BM":
-        embeddings = all_mp_net_model.encode(input.text)
+    logger.info(f"Received request with model: {input.model_name}, texts: {input.text}")
+
+    model_key = input.model_name.upper()
+    if model_key in {"BM", "BG"}:
+        embeddings = [encode_dense_cached(model_key, t) for t in input.text]
+        logger.info(f"Returning dense embeddings for {len(embeddings)} texts")
+        return {"type": "dense", "embeddings": embeddings}
+    elif model_key == "SPLADE":
+        sparse_vecs = [encode_splade_cached(t).model_dump() for t in input.text]
+        logger.info(f"Returning sparse embeddings for {len(sparse_vecs)} texts")
+        return {"type": "sparse", "embeddings": sparse_vecs}
     else:
         embeddings = bge_small_model.encode(input.text)
     return {"embeddings": embeddings.tolist()}