Skip to content

Commit

Permalink
bugfix: size msut be int type
Browse files Browse the repository at this point in the history
  • Loading branch information
wwxxzz committed Jun 7, 2024
1 parent 152e1b2 commit 2bc2d09
Show file tree
Hide file tree
Showing 2 changed files with 5 additions and 1 deletion.
2 changes: 1 addition & 1 deletion src/pai_rag/app/web/view_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -218,7 +218,7 @@ def to_app_config(self):
config["embedding"]["source"] = self.embed_source
config["embedding"]["model_name"] = self.embed_model
config["embedding"]["api_key"] = self.embed_api_key
config["embedding"]["embed_batch_size"] = self.embed_batch_size
config["embedding"]["embed_batch_size"] = int(self.embed_batch_size)

config["llm"]["source"] = self.llm
config["llm"]["endpoint"] = self.llm_eas_url
Expand Down
4 changes: 4 additions & 0 deletions src/pai_rag/modules/embedding/embedding.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,9 @@ def _create_new_instance(self, new_params: Dict[str, Any]):
source = config["source"].lower()
embed_batch_size = config.get("embed_batch_size", DEFAULT_EMBED_BATCH_SIZE)

if not isinstance(embed_batch_size, int):
raise TypeError("embed_batch_size must be of type int")

if source == "openai":
embed_model = OpenAIEmbedding(
api_key=config.get("api_key", None),
Expand All @@ -52,6 +55,7 @@ def _create_new_instance(self, new_params: Dict[str, Any]):
embed_model = HuggingFaceEmbedding(
model_name=model_path, embed_batch_size=embed_batch_size
)

logger.info(
f"Initialized HuggingFace embedding model {model_name} with {embed_batch_size} batch size."
)
Expand Down

0 comments on commit 2bc2d09

Please sign in to comment.