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

Forming embeddings-hf package

Browse source
This commit is contained in:
Cyber MacGeddon 2024-09-28 22:40:14 +01:00
parent 396a484a11
commit 3bbfe83ecf
10 changed files with 54 additions and 12 deletions
Download patch file Download diff file Expand all files Collapse all files

6
trustgraph-core/scripts/embeddings-hf
View file

@ -1,6 +0,0 @@
#!/usr/bin/env python3
from trustgraph.core.embeddings.hf import run
run()

6
trustgraph-core/setup.py
View file

@ -27,18 +27,13 @@ setuptools.setup(
python_requires='>=3.8',
download_url = "https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v" + version + ".tar.gz",
install_requires=[
"torch",
"urllib3",
"transformers",
"sentence-transformers",
"rdflib",
"pymilvus",
"langchain",
"langchain-core",
"langchain-huggingface",
"langchain-text-splitters",
"langchain-community",
"huggingface-hub",
"requests",
"cassandra-driver",
"pulsar-client",
@ -66,7 +61,6 @@ setuptools.setup(
"scripts/de-write-qdrant",
"scripts/document-rag",
"scripts/dump-parquet",
"scripts/embeddings-hf",
"scripts/embeddings-ollama",
"scripts/embeddings-vectorize",
"scripts/ge-dump-parquet",

3
trustgraph-core/trustgraph/core/embeddings/hf/__init__.py
View file

@ -1,3 +0,0 @@
from . hf import *

7
trustgraph-core/trustgraph/core/embeddings/hf/__main__.py
View file

@ -1,7 +0,0 @@
#!/usr/bin/env python3
from . hf import run
if __name__ == '__main__':
run()

99
trustgraph-core/trustgraph/core/embeddings/hf/hf.py
View file

@ -1,99 +0,0 @@
"""
Embeddings service, applies an embeddings model selected from HuggingFace.
Input is text, output is embeddings vector.
"""
from langchain_huggingface import HuggingFaceEmbeddings
from ... schema import EmbeddingsRequest, EmbeddingsResponse, Error
from ... schema import embeddings_request_queue, embeddings_response_queue
from ... log_level import LogLevel
from ... base import ConsumerProducer
module = ".".join(__name__.split(".")[1:-1])
default_input_queue = embeddings_request_queue
default_output_queue = embeddings_response_queue
default_subscriber = module
default_model="all-MiniLM-L6-v2"
class Processor(ConsumerProducer):
def __init__(self, **params):
input_queue = params.get("input_queue", default_input_queue)
output_queue = params.get("output_queue", default_output_queue)
subscriber = params.get("subscriber", default_subscriber)
model = params.get("model", default_model)
super(Processor, self).__init__(
**params | {
"input_queue": input_queue,
"output_queue": output_queue,
"subscriber": subscriber,
"input_schema": EmbeddingsRequest,
"output_schema": EmbeddingsResponse,
}
)
self.embeddings = HuggingFaceEmbeddings(model_name=model)
def handle(self, msg):
v = msg.value()
# Sender-produced ID
id = msg.properties()["id"]
print(f"Handling input {id}...", flush=True)
try:
text = v.text
embeds = self.embeddings.embed_documents([text])
print("Send response...", flush=True)
r = EmbeddingsResponse(vectors=embeds, error=None)
self.producer.send(r, properties={"id": id})
print("Done.", flush=True)
except Exception as e:
print(f"Exception: {e}")
print("Send error response...", flush=True)
r = EmbeddingsResponse(
error=Error(
type = "llm-error",
message = str(e),
),
response=None,
)
self.producer.send(r, properties={"id": id})
self.consumer.acknowledge(msg)
@staticmethod
def add_args(parser):
ConsumerProducer.add_args(
parser, default_input_queue, default_subscriber,
default_output_queue,
)
parser.add_argument(
'-m', '--model',
default="all-MiniLM-L6-v2",
help=f'LLM model (default: all-MiniLM-L6-v2)'
)
def run():
Processor.start(module, __doc__)