diff --git a/trustgraph-flow/scripts/embeddings-fastembed b/trustgraph-flow/scripts/embeddings-fastembed index 185eed59..e1322269 100755 --- a/trustgraph-flow/scripts/embeddings-fastembed +++ b/trustgraph-flow/scripts/embeddings-fastembed @@ -1,6 +1,6 @@ #!/usr/bin/env python3 -from trustgraph.embeddings.ollama import run +from trustgraph.embeddings.fastembed import run run() diff --git a/trustgraph-flow/setup.py b/trustgraph-flow/setup.py index 5407bb19..a6191cf5 100644 --- a/trustgraph-flow/setup.py +++ b/trustgraph-flow/setup.py @@ -79,6 +79,7 @@ setuptools.setup( "scripts/document-embeddings", "scripts/document-rag", "scripts/embeddings-ollama", + "scripts/embeddings-fastembed", "scripts/ge-query-milvus", "scripts/ge-query-pinecone", "scripts/ge-query-qdrant", diff --git a/trustgraph-flow/trustgraph/embeddings/fastembed/processor.py b/trustgraph-flow/trustgraph/embeddings/fastembed/processor.py index edd2346e..635387b8 100755 --- a/trustgraph-flow/trustgraph/embeddings/fastembed/processor.py +++ b/trustgraph-flow/trustgraph/embeddings/fastembed/processor.py @@ -52,11 +52,16 @@ class Processor(ConsumerProducer): print(f"Handling input {id}...", flush=True) text = v.text - embeds = self.embeddings.embed([text]) + vecs = self.embeddings.embed([text]) + + vecs = [ + v.tolist() + for v in vecs + ] print("Send response...", flush=True) r = EmbeddingsResponse( - vectors=embeddings, + vectors=list(vecs), error=None, ) @@ -78,12 +83,6 @@ class Processor(ConsumerProducer): help=f'Embeddings model (default: {default_model})' ) - parser.add_argument( - '-r', '--ollama', - default=default_ollama, - help=f'ollama (default: {default_ollama})' - ) - def run(): Processor.start(module, __doc__) diff --git a/trustgraph-flow/trustgraph/embeddings/ollama/processor.py b/trustgraph-flow/trustgraph/embeddings/ollama/processor.py index fc54cbb8..5baf64aa 100755 --- a/trustgraph-flow/trustgraph/embeddings/ollama/processor.py +++ b/trustgraph-flow/trustgraph/embeddings/ollama/processor.py @@ -1,6 +1,6 @@ """ -Embeddings service, applies an embeddings model selected from HuggingFace. +Embeddings service, applies an embeddings model hosted on a local Ollama. Input is text, output is embeddings vector. """