kris/ForcePilot
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

fix local model bugs

Browse Source
This commit is contained in:
Wenjie Zhang 2024-09-11 01:07:19 +08:00
parent ad2e6e9f92
commit 7ce9b5bd8a
3 changed files with 26 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
src/config/__init__.py
View File

@ -79,12 +79,12 @@ class Config(SimpleConfig):
def handle_self(self):
### handle local model
model_root_dir = os.getenv("MODEL_ROOT_DIR", "pretrained_models")
if self.model_local_paths is not None:
for model, model_rel_path in self.model_local_paths.items():
# 如果 model_rel_path 不是绝对路径，那么拼接 model_root_dir
if not model_rel_path.startswith("/"):
self.model_local_paths[model] = os.path.join(model_root_dir, model_rel_path)
# model_root_dir = os.getenv("MODEL_ROOT_DIR", None)
# if self.model_local_paths is not None:
# for model, model_rel_path in self.model_local_paths.items():
# # 如果 model_rel_path 不是绝对路径，那么拼接 model_root_dir
# if not model_rel_path.startswith("/"):
# self.model_local_paths[model] = os.path.join(model_root_dir, model_rel_path)
self.model_names = MODEL_NAMES

8
src/models/README.md
View File

@ -39,16 +39,14 @@ python -m vllm.entrypoints.openai.api_server \
|`zhipu`|`embedding-2`|`ZHIPUAPI` (`.env`)|
例如（`config/base.yaml`）：
例如（`saves/config/config.yaml`）：
```yaml
model_provider: qianfan
model_name: null # for default
## model dir 可以写相对路径和绝对路径
### 相对路径是相对于环境变量 (.env) 中 MODEL_ROOT_DIR 的路径
## model dir 可以写**相对路径**和**绝对路径**
### 相对路径是相对于环境变量 (.env) 中 MODEL_ROOT_DIR（若为空则是相对于 `pretrained_models`) 的路径
model_local_paths:
bge-large-zh-v1.5: /models/bge-large-zh-v1.5
```

20
src/models/embedding.py
View File

@ -15,7 +15,11 @@ GLOBAL_EMBED_STATE = {}
class EmbeddingModel(FlagModel):
def __init__(self, model_info, config, **kwargs):
self.info = model_info
model_name_or_path = config.model_local_paths.get(model_info.name, model_info.default_path)
model_name_or_path = handle_local_model(
paths=config.model_local_paths,
model_name=model_info.name,
default_path=model_info.default_path)
logger.info(f"Loading embedding model {model_info.name} from {model_name_or_path}")
super().__init__(model_name_or_path,
@ -30,7 +34,11 @@ class Reranker(FlagReranker):
assert config.reranker in RERANKER_LIST.keys(), f"Unsupported Reranker: {config.reranker}, only support {RERANKER_LIST.keys()}"
model_name_or_path = config.model_local_paths.get(config.reranker, RERANKER_LIST[config.reranker])
model_name_or_path = handle_local_model(
paths=config.model_local_paths,
model_name=config.reranker,
default_path=RERANKER_LIST[config.reranker])
logger.info(f"Loading Reranker model {config.reranker} from {model_name_or_path}")
super().__init__(model_name_or_path, use_fp16=True, **kwargs)
@ -100,4 +108,10 @@ def get_embedding_model(config):
if config.embed_model in ["zhipu-embedding-2", "zhipu-embedding-3"]:
model = ZhipuEmbedding(EMBED_MODEL_INFO[config.embed_model], config)
return model
return model
def handle_local_model(paths, model_name, default_path):
model_path = paths.get(model_name, default_path)
if os.getenv("MODEL_ROOT_DIR") and not os.path.isabs(model_path):
model_path = os.path.join(os.getenv("MODEL_ROOT_DIR"), model_path)
return model_path