apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-18 01:31:02 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

Added fastembed implementation

Browse source
This commit is contained in:
Cyber MacGeddon 2025-01-28 08:42:59 +00:00
parent 0134e0160d
commit 2dbcbbec71
1 changed files with 5 additions and 12 deletions
Download patch file Download diff file Expand all files Collapse all files

17
trustgraph-flow/trustgraph/embeddings/fastembed/processor.py
View file

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