Integration of sentence-transformer local text embedding
This commit is contained in:
@@ -157,7 +157,7 @@ class KnowledgeEmailSource:
|
||||
for uid in email_list:
|
||||
_uid = int.from_bytes(uid)
|
||||
if _uid > latest_uid:
|
||||
email_dir = self.check_email_saved(uid)
|
||||
email_dir = self.check_email_saved(uid)
|
||||
if email_dir is not None:
|
||||
logging.info(f"email uid {uid} already saved")
|
||||
else:
|
||||
|
||||
@@ -17,18 +17,27 @@ class LocalSentenceTransformer_ComputeNode(Queue_ComputeNode):
|
||||
def __init__(self, model_name: str = "all-MiniLM-L6-v2"):
|
||||
super().__init__()
|
||||
|
||||
logger.info(
|
||||
f"LocalSentenceTransformer_ComputeNode init, model_name: {model_name}"
|
||||
)
|
||||
self.node_id = "local_sentence_transformer_node"
|
||||
self.model_name = model_name
|
||||
self.model = None
|
||||
|
||||
def initial(self) -> bool:
|
||||
logger.info(
|
||||
f"LocalSentenceTransformer_ComputeNode init, model_name: {self.model_name}"
|
||||
)
|
||||
|
||||
assert self.model_name is not None
|
||||
assert self.model is None
|
||||
try:
|
||||
from sentence_transformers import SentenceTransformer
|
||||
|
||||
self.model = SentenceTransformer(self.model)
|
||||
except Exception as err:
|
||||
logger.error(f"load model {self.model} failed: {err}")
|
||||
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
async def execute_task(
|
||||
self, task: ComputeTask
|
||||
) -> {
|
||||
|
||||
Reference in New Issue
Block a user