diff --git a/trustgraph-flow/trustgraph/embeddings/fastembed/processor.py b/trustgraph-flow/trustgraph/embeddings/fastembed/processor.py index fc54cbb8..edd2346e 100755 --- a/trustgraph-flow/trustgraph/embeddings/fastembed/processor.py +++ b/trustgraph-flow/trustgraph/embeddings/fastembed/processor.py @@ -8,7 +8,7 @@ from ... schema import EmbeddingsRequest, EmbeddingsResponse from ... schema import embeddings_request_queue, embeddings_response_queue from ... log_level import LogLevel from ... base import ConsumerProducer -from ollama import Client +from fastembed import TextEmbedding import os module = ".".join(__name__.split(".")[1:-1]) @@ -16,8 +16,7 @@ module = ".".join(__name__.split(".")[1:-1]) default_input_queue = embeddings_request_queue default_output_queue = embeddings_response_queue default_subscriber = module -default_model="mxbai-embed-large" -default_ollama = os.getenv("OLLAMA_HOST", 'http://localhost:11434') +default_model="sentence-transformers/all-MiniLM-L6-v2" class Processor(ConsumerProducer): @@ -27,7 +26,6 @@ class Processor(ConsumerProducer): output_queue = params.get("output_queue", default_output_queue) subscriber = params.get("subscriber", default_subscriber) - ollama = params.get("ollama", default_ollama) model = params.get("model", default_model) super(Processor, self).__init__( @@ -37,13 +35,11 @@ class Processor(ConsumerProducer): "subscriber": subscriber, "input_schema": EmbeddingsRequest, "output_schema": EmbeddingsResponse, - "ollama": ollama, "model": model, } ) - self.client = Client(host=ollama) - self.model = model + self.embeddings = TextEmbedding(model_name = model) def handle(self, msg): @@ -56,14 +52,11 @@ class Processor(ConsumerProducer): print(f"Handling input {id}...", flush=True) text = v.text - embeds = self.client.embed( - model = self.model, - input = text - ) + embeds = self.embeddings.embed([text]) print("Send response...", flush=True) r = EmbeddingsResponse( - vectors=embeds.embeddings, + vectors=embeddings, error=None, )