Update src/rag_engine.py

src/rag_engine.py

@@ -2,6 +2,7 @@ import os
 import shutil
 import logging
 from typing import List, Tuple, Optional
+from huggingface_hub import snapshot_download
 from langchain_community.document_loaders import PyPDFLoader, TextLoader, UnstructuredWordDocumentLoader, UnstructuredPowerPointLoader
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_huggingface import HuggingFaceEmbeddings
@@ -27,15 +28,40 @@ except Exception as e:
 
 def get_embedding_func(model_name: str = "sentence-transformers/all-MiniLM-L6-v2"):
     try:
-        if "openai" in model_name.lower():
+        # CHECK 1: OpenAI
+        if "openai" in model_name.lower() or "text-embedding" in model_name.lower():
             if not os.getenv("OPENAI_API_KEY"): raise ValueError("OpenAI API Key not found.")
             return OpenAIEmbeddings(model=model_name)
-
+            
+        # CHECK 2: YOUR CUSTOM FINE-TUNE (Updated for Subfolders)
         elif "navy-custom-models" in model_name:
-            return HuggingFaceEmbeddings(model_name=model_name)
-        
+             logger.info(f"Downloading custom model from: {model_name}")
+             
+             # 1. Parse the repo and folder from your string
+             # Input: "NavyDevilDoc/navy-custom-models/bge-finetuned"
+             parts = model_name.split("/")
+             # Repo ID is the first two parts: "NavyDevilDoc/navy-custom-models"
+             repo_id = f"{parts[0]}/{parts[1]}"
+             # Folder is the rest: "bge-finetuned"
+             folder_name = parts[2]
+             
+             # 2. Download ONLY that folder
+             storage_path = snapshot_download(
+                 repo_id=repo_id, 
+                 repo_type="model", 
+                 allow_patterns=f"{folder_name}/*"
+             )
+             
+             # 3. Construct the local path to the inner folder
+             local_model_path = os.path.join(storage_path, folder_name)
+             
+             # 4. Load from the local path
+             return HuggingFaceEmbeddings(model_name=local_model_path)
+             
+        # CHECK 3: Standard Public Models
         else:
             return HuggingFaceEmbeddings(model_name=model_name)
+            
     except Exception as e:
         logger.error(f"Failed to load embedding model '{model_name}': {e}")
         return HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")