diff --git a/src/aios_kernel/local_st_compute_node.py b/src/aios_kernel/local_st_compute_node.py index bc9f4af..e73b9f9 100644 --- a/src/aios_kernel/local_st_compute_node.py +++ b/src/aios_kernel/local_st_compute_node.py @@ -89,7 +89,7 @@ class LocalSentenceTransformer_ComputeNode(Queue_ComputeNode): def is_support(self, task: ComputeTask) -> bool: return task.task_type == ComputeTaskType.TEXT_EMBEDDING and ( - not task.params["model_name"] or task.params["model_name"] == "all-MiniLM-L6-v2" + not task.params["model_name"] or task.params["model_name"] == "llama" ) def is_local(self) -> bool: diff --git a/src/requirements.txt b/src/requirements.txt index 3f19c10..ec9ef4f 100644 --- a/src/requirements.txt +++ b/src/requirements.txt @@ -136,4 +136,4 @@ aiosqlite python-telegram-bot pydub stability_sdk -sentence-transformers==2.2 +sentence-transformers==2.2.2