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

fix model info bugs

Browse Source
This commit is contained in:
Wenjie Zhang 2024-11-22 00:03:25 +08:00
parent 8c26722aa9
commit 6433046c0b
1 changed files with 5 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
src/models/embedding.py
View File

@ -17,16 +17,16 @@ class EmbeddingModel(FlagModel):
self.info = model_info self.info = model_info
model_name_or_path = handle_local_model( model_name_or_path = handle_local_model(
paths=config.model_local_paths, paths=config.model_local_paths,
model_name=model_info.name, model_name=model_info["name"],
default_path=model_info.default_path) default_path=model_info.get("default_path", None))
logger.info(f"Loading embedding model {model_info.name} from {model_name_or_path}") logger.info(f"Loading embedding model {model_info["name"]} from {model_name_or_path}")
super().__init__(model_name_or_path, super().__init__(model_name_or_path,
query_instruction_for_retrieval=model_info.get("query_instruction", None), query_instruction_for_retrieval=model_info.get("query_instruction", None),
use_fp16=False, **kwargs) use_fp16=False, **kwargs)
logger.info(f"Embedding model {model_info.name} loaded") logger.info(f"Embedding model {model_info["name"]} loaded")
class Reranker(FlagReranker): class Reranker(FlagReranker):
@ -77,7 +77,7 @@ class ZhipuEmbedding:
group_msg = message[i:i+batch_size] group_msg = message[i:i+batch_size]
response = self.client.embeddings.create( response = self.client.embeddings.create(
model=self.model_info.default_path, model=self.model_info.get("default_path", None),
input=group_msg, input=group_msg,
) )