Update embedding model to Vietnamese_Embedding and adjust related configurations. Replace FlagEmbedding with SentenceTransformer in embedding service, and ensure dimension checks for vectors in Pinecone service. Update requirements to reflect new dependencies.

app.py

@@ -19,7 +19,7 @@ async def lifespan(app: FastAPI):
 
 app = FastAPI(
     title="MentorMe AI Recommendation Server",
-    description="AI-powered mentor-mentee recommendation using bge-m3 embeddings and Pinecone",
+    description="AI-powered mentor-mentee recommendation using Vietnamese_Embedding and Pinecone",
     version="1.0.0",
     lifespan=lifespan
 )

config/settings.py

@@ -22,7 +22,7 @@ class Settings(BaseSettings):
     PORT: int = int(os.getenv("PORT", "7860"))
     DEBUG: bool = False
     
-    BGE_MODEL_NAME: str = "BAAI/bge-m3"
+    EMBEDDING_MODEL_NAME: str = "AITeamVN/Vietnamese_Embedding"
     USE_FP16: bool = True
     
     RECOMMENDATION_TOP_K: int = 30

main.py

@@ -19,7 +19,7 @@ async def lifespan(app: FastAPI):
 
 app = FastAPI(
     title="MentorMe AI Recommendation Server",
-    description="AI-powered mentor-mentee recommendation using bge-m3 embeddings and Pinecone",
+    description="AI-powered mentor-mentee recommendation using Vietnamese_Embedding and Pinecone",
     version="1.0.0",
     lifespan=lifespan
 )

requirements.txt

@@ -4,12 +4,11 @@ pydantic==2.5.0
 pydantic-settings==2.1.0
 python-dotenv==1.0.0
 pinecone-client>=3.2.0,<6.0.0
-FlagEmbedding==1.2.11
+sentence-transformers>=2.2.0
 torch>=2.0.0
 numpy>=1.24.0
 python-multipart==0.0.6
 transformers>=4.30.0
-peft>=0.3.0
 accelerate>=0.20.0
 sentencepiece>=0.1.99
 

services/embedding_service.py

@@ -1,5 +1,5 @@
 import torch
-from FlagEmbedding import FlagModel
+from sentence_transformers import SentenceTransformer
 import logging
 from typing import List, Union
 from config.settings import get_settings
@@ -22,18 +22,16 @@ class EmbeddingService:
     def _load_model(self):
         settings = get_settings()
         try:
-            logger.info(f"Loading embedding model: {settings.BGE_MODEL_NAME}")
+            logger.info(f"Loading embedding model: {settings.EMBEDDING_MODEL_NAME}")
             device = "cuda" if torch.cuda.is_available() else "cpu"
             logger.info(f"Using device: {device}")
             
-            EmbeddingService._model = FlagModel(
-                settings.BGE_MODEL_NAME,
-                use_fp16=settings.USE_FP16 and device == "cuda",
-                query_instruction_for_retrieval="Represent this sentence for searching relevant passages:"
+            EmbeddingService._model = SentenceTransformer(
+                settings.EMBEDDING_MODEL_NAME,
+                device=device
             )
             
-            if device == "cuda":
-                EmbeddingService._model = EmbeddingService._model.cuda()
+            EmbeddingService._model.max_seq_length = 2048
             
             logger.info("Embedding model loaded successfully")
         except Exception as e:
@@ -45,7 +43,7 @@ class EmbeddingService:
         texts: Union[str, List[str]],
         is_query: bool = False,
         batch_size: int = 32,
-        max_length: int = 8192
+        max_length: int = 2048
     ) -> Union[List[float], List[List[float]]]:
         if EmbeddingService._model is None:
             raise RuntimeError("Embedding model not loaded")
@@ -56,34 +54,53 @@ class EmbeddingService:
         else:
             single_text = False
         
+        if not texts:
+            raise ValueError("Texts cannot be empty")
+        
         try:
-            if is_query:
-                embeddings = EmbeddingService._model.encode_queries(
-                    texts,
-                    batch_size=batch_size,
-                    max_length=max_length
-                )
-            else:
-                embeddings = EmbeddingService._model.encode(
-                    texts,
-                    batch_size=batch_size,
-                    max_length=max_length
-                )
+            embeddings = EmbeddingService._model.encode(
+                texts,
+                batch_size=batch_size,
+                show_progress_bar=False,
+                convert_to_numpy=True,
+                normalize_embeddings=False
+            )
             
+            expected_dim = 1024
             if single_text:
-                return embeddings[0].tolist()
-            return [emb.tolist() for emb in embeddings]
+                embedding_list = embeddings[0].tolist()
+                if len(embedding_list) != expected_dim:
+                    logger.warning(f"Embedding dimension mismatch: expected {expected_dim}, got {len(embedding_list)}")
+                return embedding_list
+            
+            result = []
+            for emb in embeddings:
+                emb_list = emb.tolist()
+                if len(emb_list) != expected_dim:
+                    logger.warning(f"Embedding dimension mismatch: expected {expected_dim}, got {len(emb_list)}")
+                result.append(emb_list)
+            
+            return result
         except Exception as e:
             logger.error(f"Error encoding texts: {str(e)}")
             raise
     
     def get_model_info(self) -> dict:
         settings = get_settings()
+        dimension = 1024
+        
+        if EmbeddingService._model is not None:
+            try:
+                test_embedding = EmbeddingService._model.encode(["test"], convert_to_numpy=True)
+                dimension = len(test_embedding[0])
+            except Exception as e:
+                logger.warning(f"Could not determine model dimension: {str(e)}")
+        
         return {
-            "model_name": settings.BGE_MODEL_NAME,
-            "dimension": 1024,
+            "model_name": settings.EMBEDDING_MODEL_NAME,
+            "dimension": dimension,
             "device": "cuda" if torch.cuda.is_available() else "cpu",
-            "fp16": settings.USE_FP16 and torch.cuda.is_available()
+            "max_seq_length": EmbeddingService._model.max_seq_length if EmbeddingService._model else 2048
         }
 
 

services/pinecone_service.py

@@ -58,6 +58,14 @@ class PineconeService:
         metadata: Dict[str, Any]
     ) -> bool:
         try:
+            settings = get_settings()
+            expected_dim = settings.PINECONE_DIMENSION
+            
+            if len(vector) != expected_dim:
+                error_msg = f"Vector dimension mismatch: expected {expected_dim}, got {len(vector)}"
+                logger.error(error_msg)
+                raise ValueError(error_msg)
+            
             PineconeService._index.upsert(
                 vectors=[{
                     "id": str(mentor_id),
@@ -91,6 +99,14 @@ class PineconeService:
         include_metadata: bool = True
     ) -> List[Dict[str, Any]]:
         try:
+            settings = get_settings()
+            expected_dim = settings.PINECONE_DIMENSION
+            
+            if len(query_vector) != expected_dim:
+                error_msg = f"Query vector dimension mismatch: expected {expected_dim}, got {len(query_vector)}"
+                logger.error(error_msg)
+                raise ValueError(error_msg)
+            
             query_response = PineconeService._index.query(
                 vector=query_vector,
                 top_k=top_k,