diff --git a/src/models/embedding.py b/src/models/embedding.py index 0cd9c5fb..8a28f237 100644 --- a/src/models/embedding.py +++ b/src/models/embedding.py @@ -17,16 +17,16 @@ class EmbeddingModel(FlagModel): self.info = model_info model_name_or_path = handle_local_model( paths=config.model_local_paths, - model_name=model_info.name, - default_path=model_info.default_path) + model_name=model_info["name"], + 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, query_instruction_for_retrieval=model_info.get("query_instruction", None), 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): @@ -77,7 +77,7 @@ class ZhipuEmbedding: group_msg = message[i:i+batch_size] response = self.client.embeddings.create( - model=self.model_info.default_path, + model=self.model_info.get("default_path", None), input=group_msg, )