diff --git a/src/aios_kernel/__init__.py b/src/aios_kernel/__init__.py index debd524..ee0b02d 100644 --- a/src/aios_kernel/__init__.py +++ b/src/aios_kernel/__init__.py @@ -25,5 +25,6 @@ from .local_stability_node import Local_Stability_ComputeNode from .stability_node import Stability_ComputeNode from .local_st_compute_node import LocalSentenceTransformer_Text_ComputeNode,LocalSentenceTransformer_Image_ComputeNode + AIOS_Version = "0.5.1, build 2023-9-27" diff --git a/src/aios_kernel/local_st_compute_node.py b/src/aios_kernel/local_st_compute_node.py index 43c2114..cfef1e0 100644 --- a/src/aios_kernel/local_st_compute_node.py +++ b/src/aios_kernel/local_st_compute_node.py @@ -58,7 +58,7 @@ class LocalSentenceTransformer_Text_ComputeNode(Queue_ComputeNode): logger.debug( f"LocalSentenceTransformer_Text_ComputeNode task input: {input}" ) - sentence_embeddings = self.model.encode(input).tolist() + sentence_embeddings = self.model.encode(input, show_progress_bar=False).tolist() # logger.debug(f"LocalSentenceTransformer_Text_ComputeNode task sentence_embeddings: {sentence_embeddings}") return { "state": ComputeTaskState.DONE, @@ -87,7 +87,7 @@ class LocalSentenceTransformer_Text_ComputeNode(Queue_ComputeNode): pass def is_support(self, task: ComputeTask) -> bool: - return task.task_type == ComputeTaskType.TEXT_EMBEDDING + return task.task_type == ComputeTaskType.TEXT_EMBEDDING and task.params["model_name"] == "all-MiniLM-L6-v2" def is_local(self) -> bool: return True @@ -187,7 +187,7 @@ class LocalSentenceTransformer_Image_ComputeNode(Queue_ComputeNode): logger.debug( f"LocalSentenceTransformer_Text_ComputeNode task text input: {input}" ) - sentence_embeddings = self.multi_model.encode(input) + sentence_embeddings = self.multi_model.encode(input, show_progress_bar=False).tolist() # logger.debug(f"LocalSentenceTransformer_Text_ComputeNode task sentence_embeddings: {sentence_embeddings}") return { "state": ComputeTaskState.DONE, @@ -238,7 +238,7 @@ class LocalSentenceTransformer_Image_ComputeNode(Queue_ComputeNode): def is_support(self, task: ComputeTask) -> bool: return ( - task.task_type == ComputeTaskType.TEXT_EMBEDDING + (task.task_type == ComputeTaskType.TEXT_EMBEDDING and task.params["model_name"] == "clip-ViT-B-32") or task.task_type == ComputeTaskType.IMAGE_EMBEDDING ) diff --git a/test/test_embedding.py b/test/test_embedding.py index 5b14dda..7141bf5 100644 --- a/test/test_embedding.py +++ b/test/test_embedding.py @@ -67,11 +67,13 @@ def test_st(): ] # Compute embeddings + #embeddings = model.encode(sentences, convert_to_tensor=True) embeddings = model.encode(sentences) print("embeddings as follows: ") print(embeddings) + # Compute cosine-similarities for each sentence with each other sentence cosine_scores = util.cos_sim(embeddings, embeddings)