From 8696bf65837bd490e089e76d4082c14aeaf7ecff Mon Sep 17 00:00:00 2001
From: Wenjie Zhang <xerrors@163.com>
Date: Sun, 13 Apr 2025 23:50:19 +0800
Subject: [PATCH] add logger for loading local model

---
 src/config/__init__.py  |  8 +++++++-
 src/models/embedding.py | 26 +++++++++++++++++++++++---
 2 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/src/config/__init__.py b/src/config/__init__.py
index b8cf6119..d72508ab 100644
--- a/src/config/__init__.py
+++ b/src/config/__init__.py
@@ -121,7 +121,13 @@ class Config(SimpleConfig):
         """
         model_provider_info = self.model_names.get(self.model_provider, {})
         self.model_dir = os.environ.get("MODEL_DIR", "")
-        logger.info(f"MODEL_DIR: {self.model_dir}; 如果是在 docker 中运行，会自动挂载 MODEL_DIR 到 /models 目录，请检查 docker compose 文件")
+
+        if self.model_dir:
+            if os.path.exists(self.model_dir):
+                logger.info(f"MODEL_DIR （{self.model_dir}） 下面的文件夹: {os.listdir(self.model_dir)}")
+            else:
+                logger.warning(f"MODEL_DIR （{self.model_dir}） 不存在，请检查是否配置正确，比如 docker-compose 文件中的映射")
+
 
         # 检查模型提供商是否存在
         if self.model_provider != "custom":
diff --git a/src/models/embedding.py b/src/models/embedding.py
index aa035bc8..7ea83dd6 100644
--- a/src/models/embedding.py
+++ b/src/models/embedding.py
@@ -53,6 +53,19 @@ class BaseEmbeddingModel:
 
 class LocalEmbeddingModel(FlagModel, BaseEmbeddingModel):
     def __init__(self, config, **kwargs):
+        """
+        对于本地模型，也可以在 src/static/models.private.yaml 中配置对应的 local_path 路径
+
+        ```yaml
+        EMBED_MODEL_INFO:
+            local/BAAI/bge-m3:
+                dimension: 1024
+                name: BAAI/bge-m3
+                local_path: /path/to/bge-m3
+        ```
+
+        但是也要确保在 docker-compose 中映射了 MODEL_DIR 到 /models 目录
+        """
         info = config.embed_model_names[config.embed_model]
 
         self.model = config.model_local_paths.get(info["name"], info.get("local_path"))
@@ -60,10 +73,17 @@ class LocalEmbeddingModel(FlagModel, BaseEmbeddingModel):
         self.dimension = info["dimension"]
         self.embed_model_fullname = config.embed_model
 
-        if os.path.exists(_path := os.path.join(os.getenv("MODEL_DIR"), self.model)):
-            self.model = _path
+        if os.getenv("MODEL_DIR"):
+            if os.path.exists(_path := os.path.join(os.getenv("MODEL_DIR"), self.model)):
+                self.model = _path
+            else:
+                logger.warning(f"Local model `{info['name']}` not found in `{self.model}`, using `{info['name']}`")
 
-        logger.info(f"Loading local model `{info['name']}` from `{self.model}` with device `{config.device}`")
+        logger.info(f"Loading local model `{info['name']}` from `{self.model}` with device `{config.device}`，"
+                    f"如果没配置任何路径的话，正常情况下会自动从 Huggingface 下载模型，如果遇到下载失败，可以尝试使用 HF_MIRROR 环境变量；"
+                    f"如果还是不行，建议手动下载到某个文件夹比如  /path/to/models/BAAI/bge-m3 目录下；"
+                    f"然后配置 src/.env 文件中的 MODEL_DIR 环境变量到 /path/to/models 目录；"
+                    f"如果是在 docker 中运行，请确保 docker-compose 文件（line 12 左右）中映射了 MODEL_DIR 到 /models 目录")
 
         super().__init__(self.model,
                 query_instruction_for_retrieval=info.get("query_instruction", None),