mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-10 05:42:12 +02:00
Feature/refactor entity embeddings (#235)
* Make schema changes * Core entity context flow in place * extract-def outputs entity contexts * Refactored qdrant write * Refactoring of all vector stores in place
This commit is contained in:
cybermaggedon
GitHub
parent
9942f63773
commit
a458d57af2
9 changed files with 230 additions and 169 deletions
|
|
@ -1,14 +1,17 @@
|
|||
|
||||
"""
|
||||
Simple decoder, accepts embeddings+text chunks input, applies entity analysis to
|
||||
get entity definitions which are output as graph edges.
|
||||
Simple decoder, accepts text chunks input, applies entity analysis to
|
||||
get entity definitions which are output as graph edges along with
|
||||
entity/context definitions for embedding.
|
||||
"""
|
||||
|
||||
import urllib.parse
|
||||
import json
|
||||
from pulsar.schema import JsonSchema
|
||||
|
||||
from .... schema import ChunkEmbeddings, Triple, Triples, Metadata, Value
|
||||
from .... schema import chunk_embeddings_ingest_queue, triples_store_queue
|
||||
from .... schema import Chunk, Triple, Triples, Metadata, Value
|
||||
from .... schema import EntityContext, EntityContexts
|
||||
from .... schema import chunk_ingest_queue, triples_store_queue
|
||||
from .... schema import entity_contexts_ingest_queue
|
||||
from .... schema import prompt_request_queue
|
||||
from .... schema import prompt_response_queue
|
||||
from .... log_level import LogLevel
|
||||
|
|
@ -22,8 +25,9 @@ SUBJECT_OF_VALUE = Value(value=SUBJECT_OF, is_uri=True)
|
|||
|
||||
module = ".".join(__name__.split(".")[1:-1])
|
||||
|
||||
default_input_queue = chunk_embeddings_ingest_queue
|
||||
default_input_queue = chunk_ingest_queue
|
||||
default_output_queue = triples_store_queue
|
||||
default_entity_context_queue = entity_contexts_ingest_queue
|
||||
default_subscriber = module
|
||||
|
||||
class Processor(ConsumerProducer):
|
||||
|
|
@ -32,6 +36,10 @@ class Processor(ConsumerProducer):
|
|||
|
||||
input_queue = params.get("input_queue", default_input_queue)
|
||||
output_queue = params.get("output_queue", default_output_queue)
|
||||
ec_queue = params.get(
|
||||
"entity_context_queue",
|
||||
default_entity_context_queue
|
||||
)
|
||||
subscriber = params.get("subscriber", default_subscriber)
|
||||
pr_request_queue = params.get(
|
||||
"prompt_request_queue", prompt_request_queue
|
||||
|
|
@ -45,13 +53,30 @@ class Processor(ConsumerProducer):
|
|||
"input_queue": input_queue,
|
||||
"output_queue": output_queue,
|
||||
"subscriber": subscriber,
|
||||
"input_schema": ChunkEmbeddings,
|
||||
"input_schema": Chunk,
|
||||
"output_schema": Triples,
|
||||
"prompt_request_queue": pr_request_queue,
|
||||
"prompt_response_queue": pr_response_queue,
|
||||
}
|
||||
)
|
||||
|
||||
self.ec_prod = self.client.create_producer(
|
||||
topic=ec_queue,
|
||||
schema=JsonSchema(EntityContexts),
|
||||
)
|
||||
|
||||
__class__.pubsub_metric.info({
|
||||
"input_queue": input_queue,
|
||||
"output_queue": output_queue,
|
||||
"entity_context_queue": ec_queue,
|
||||
"prompt_request_queue": pr_request_queue,
|
||||
"prompt_response_queue": pr_response_queue,
|
||||
"subscriber": subscriber,
|
||||
"input_schema": Chunk.__name__,
|
||||
"output_schema": Triples.__name__,
|
||||
"vector_schema": EntityContexts.__name__,
|
||||
})
|
||||
|
||||
self.prompt = PromptClient(
|
||||
pulsar_host=self.pulsar_host,
|
||||
input_queue=pr_request_queue,
|
||||
|
|
@ -79,6 +104,14 @@ class Processor(ConsumerProducer):
|
|||
)
|
||||
self.producer.send(t)
|
||||
|
||||
def emit_ecs(self, metadata, entities):
|
||||
|
||||
t = EntityContexts(
|
||||
metadata=metadata,
|
||||
entities=entities,
|
||||
)
|
||||
self.ec_prod.send(t)
|
||||
|
||||
def handle(self, msg):
|
||||
|
||||
v = msg.value()
|
||||
|
|
@ -91,6 +124,7 @@ class Processor(ConsumerProducer):
|
|||
defs = self.get_definitions(chunk)
|
||||
|
||||
triples = []
|
||||
entities = []
|
||||
|
||||
# FIXME: Putting metadata into triples store is duplicated in
|
||||
# relationships extractor too
|
||||
|
|
@ -129,6 +163,14 @@ class Processor(ConsumerProducer):
|
|||
o=Value(value=v.metadata.id, is_uri=True)
|
||||
))
|
||||
|
||||
ec = EntityContext(
|
||||
entity=s_value,
|
||||
context=defn.definition,
|
||||
)
|
||||
|
||||
entities.append(ec)
|
||||
|
||||
|
||||
self.emit_edges(
|
||||
Metadata(
|
||||
id=v.metadata.id,
|
||||
|
|
@ -139,6 +181,16 @@ class Processor(ConsumerProducer):
|
|||
triples
|
||||
)
|
||||
|
||||
self.emit_ecs(
|
||||
Metadata(
|
||||
id=v.metadata.id,
|
||||
metadata=[],
|
||||
user=v.metadata.user,
|
||||
collection=v.metadata.collection,
|
||||
),
|
||||
entities
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
print("Exception: ", e, flush=True)
|
||||
|
||||
|
|
@ -152,6 +204,12 @@ class Processor(ConsumerProducer):
|
|||
default_output_queue,
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'-e', '--entity-context-queue',
|
||||
default=default_entity_context_queue,
|
||||
help=f'Entity context queue (default: {default_entity_context_queue})'
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--prompt-request-queue',
|
||||
default=prompt_request_queue,
|
||||
|
|
|
|||
|
|
@ -1,18 +1,15 @@
|
|||
|
||||
"""
|
||||
Simple decoder, accepts vector+text chunks input, applies entity
|
||||
Simple decoder, accepts text chunks input, applies entity
|
||||
relationship analysis to get entity relationship edges which are output as
|
||||
graph edges.
|
||||
"""
|
||||
|
||||
import urllib.parse
|
||||
import os
|
||||
from pulsar.schema import JsonSchema
|
||||
|
||||
from .... schema import ChunkEmbeddings, Triple, Triples, GraphEmbeddings
|
||||
from .... schema import Chunk, Triple, Triples
|
||||
from .... schema import Metadata, Value
|
||||
from .... schema import chunk_embeddings_ingest_queue, triples_store_queue
|
||||
from .... schema import graph_embeddings_store_queue
|
||||
from .... schema import chunk_ingest_queue, triples_store_queue
|
||||
from .... schema import prompt_request_queue
|
||||
from .... schema import prompt_response_queue
|
||||
from .... log_level import LogLevel
|
||||
|
|
@ -25,9 +22,8 @@ SUBJECT_OF_VALUE = Value(value=SUBJECT_OF, is_uri=True)
|
|||
|
||||
module = ".".join(__name__.split(".")[1:-1])
|
||||
|
||||
default_input_queue = chunk_embeddings_ingest_queue
|
||||
default_input_queue = chunk_ingest_queue
|
||||
default_output_queue = triples_store_queue
|
||||
default_vector_queue = graph_embeddings_store_queue
|
||||
default_subscriber = module
|
||||
|
||||
class Processor(ConsumerProducer):
|
||||
|
|
@ -36,7 +32,6 @@ class Processor(ConsumerProducer):
|
|||
|
||||
input_queue = params.get("input_queue", default_input_queue)
|
||||
output_queue = params.get("output_queue", default_output_queue)
|
||||
vector_queue = params.get("vector_queue", default_vector_queue)
|
||||
subscriber = params.get("subscriber", default_subscriber)
|
||||
pr_request_queue = params.get(
|
||||
"prompt_request_queue", prompt_request_queue
|
||||
|
|
@ -50,30 +45,13 @@ class Processor(ConsumerProducer):
|
|||
"input_queue": input_queue,
|
||||
"output_queue": output_queue,
|
||||
"subscriber": subscriber,
|
||||
"input_schema": ChunkEmbeddings,
|
||||
"input_schema": Chunk,
|
||||
"output_schema": Triples,
|
||||
"prompt_request_queue": pr_request_queue,
|
||||
"prompt_response_queue": pr_response_queue,
|
||||
}
|
||||
)
|
||||
|
||||
self.vec_prod = self.client.create_producer(
|
||||
topic=vector_queue,
|
||||
schema=JsonSchema(GraphEmbeddings),
|
||||
)
|
||||
|
||||
__class__.pubsub_metric.info({
|
||||
"input_queue": input_queue,
|
||||
"output_queue": output_queue,
|
||||
"vector_queue": vector_queue,
|
||||
"prompt_request_queue": pr_request_queue,
|
||||
"prompt_response_queue": pr_response_queue,
|
||||
"subscriber": subscriber,
|
||||
"input_schema": ChunkEmbeddings.__name__,
|
||||
"output_schema": Triples.__name__,
|
||||
"vector_schema": GraphEmbeddings.__name__,
|
||||
})
|
||||
|
||||
self.prompt = PromptClient(
|
||||
pulsar_host=self.pulsar_host,
|
||||
input_queue=pr_request_queue,
|
||||
|
|
@ -101,11 +79,6 @@ class Processor(ConsumerProducer):
|
|||
)
|
||||
self.producer.send(t)
|
||||
|
||||
def emit_vec(self, metadata, ent, vec):
|
||||
|
||||
r = GraphEmbeddings(metadata=metadata, entity=ent, vectors=vec)
|
||||
self.vec_prod.send(r)
|
||||
|
||||
def handle(self, msg):
|
||||
|
||||
v = msg.value()
|
||||
|
|
@ -193,12 +166,6 @@ class Processor(ConsumerProducer):
|
|||
o=Value(value=v.metadata.id, is_uri=True)
|
||||
))
|
||||
|
||||
self.emit_vec(v.metadata, s_value, v.vectors)
|
||||
self.emit_vec(v.metadata, p_value, v.vectors)
|
||||
|
||||
if rel.o_entity:
|
||||
self.emit_vec(v.metadata, o_value, v.vectors)
|
||||
|
||||
self.emit_edges(
|
||||
Metadata(
|
||||
id=v.metadata.id,
|
||||
|
|
@ -222,12 +189,6 @@ class Processor(ConsumerProducer):
|
|||
default_output_queue,
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'-c', '--vector-queue',
|
||||
default=default_vector_queue,
|
||||
help=f'Vector output queue (default: {default_vector_queue})'
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--prompt-request-queue',
|
||||
default=prompt_request_queue,
|
||||
|
|
|
|||
|
|
@ -1,14 +1,14 @@
|
|||
|
||||
"""
|
||||
Simple decoder, accepts embeddings+text chunks input, applies entity analysis to
|
||||
get entity definitions which are output as graph edges.
|
||||
Simple decoder, accepts text chunks input, applies entity analysis to
|
||||
get topics which are output as graph edges.
|
||||
"""
|
||||
|
||||
import urllib.parse
|
||||
import json
|
||||
|
||||
from .... schema import ChunkEmbeddings, Triple, Triples, Metadata, Value
|
||||
from .... schema import chunk_embeddings_ingest_queue, triples_store_queue
|
||||
from .... schema import Chunk, Triple, Triples, Metadata, Value
|
||||
from .... schema import chunk_ingest_queue, triples_store_queue
|
||||
from .... schema import prompt_request_queue
|
||||
from .... schema import prompt_response_queue
|
||||
from .... log_level import LogLevel
|
||||
|
|
@ -20,7 +20,7 @@ DEFINITION_VALUE = Value(value=DEFINITION, is_uri=True)
|
|||
|
||||
module = ".".join(__name__.split(".")[1:-1])
|
||||
|
||||
default_input_queue = chunk_embeddings_ingest_queue
|
||||
default_input_queue = chunk_ingest_queue
|
||||
default_output_queue = triples_store_queue
|
||||
default_subscriber = module
|
||||
|
||||
|
|
@ -43,7 +43,7 @@ class Processor(ConsumerProducer):
|
|||
"input_queue": input_queue,
|
||||
"output_queue": output_queue,
|
||||
"subscriber": subscriber,
|
||||
"input_schema": ChunkEmbeddings,
|
||||
"input_schema": Chunk,
|
||||
"output_schema": Triples,
|
||||
"prompt_request_queue": pr_request_queue,
|
||||
"prompt_response_queue": pr_response_queue,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue