From 5f346b1bd9f5997a252fceda48a6a447cb5a81e7 Mon Sep 17 00:00:00 2001 From: liyaxing Date: Mon, 25 Sep 2023 16:22:15 +0800 Subject: [PATCH] Add sentence-transformer local text embedding supports --- src/aios_kernel/local_st_compute_node.py | 2 +- src/requirements.txt | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) 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