trustgraph-span 2.11.1__tar.gz

trustgraph_span-2.11.1/PKG-INFO

@@ -0,0 +1,19 @@
+Metadata-Version: 2.4
+Name: trustgraph-span
+Version: 2.11.1
+Summary: Span-extraction field extraction service for TrustGraph.
+Author-email: "trustgraph.ai" <security@trustgraph.ai>
+Project-URL: Homepage, https://github.com/trustgraph-ai/trustgraph
+Classifier: Programming Language :: Python :: 3
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: trustgraph-base<2.12,>=2.11
+Requires-Dist: trustgraph-flow<2.12,>=2.11
+Requires-Dist: onnxruntime
+Requires-Dist: tokenizers
+Requires-Dist: huggingface-hub
+Requires-Dist: numpy
+Requires-Dist: pulsar-client
+Requires-Dist: pyyaml
+Requires-Dist: prometheus-client

trustgraph_span-2.11.1/pyproject.toml

@@ -0,0 +1,38 @@
+[build-system]
+requires = ["setuptools>=61.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "trustgraph-span"
+dynamic = ["version"]
+authors = [{name = "trustgraph.ai", email = "security@trustgraph.ai"}]
+description = "Span-extraction field extraction service for TrustGraph."
+readme = "README.md"
+requires-python = ">=3.10"
+dependencies = [
+    "trustgraph-base>=2.11,<2.12",
+    "trustgraph-flow>=2.11,<2.12",
+    "onnxruntime",
+    "tokenizers",
+    "huggingface-hub",
+    "numpy",
+    "pulsar-client",
+    "pyyaml",
+    "prometheus-client",
+]
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "Operating System :: OS Independent",
+]
+
+[project.urls]
+Homepage = "https://github.com/trustgraph-ai/trustgraph"
+
+[project.scripts]
+field-extraction-span = "trustgraph.model.field_extraction.span:run"
+
+[tool.setuptools.packages.find]
+include = ["trustgraph*"]
+
+[tool.setuptools.dynamic]
+version = {attr = "trustgraph.span_version.__version__"}

trustgraph_span-2.11.1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

trustgraph_span-2.11.1/trustgraph/__init__.py

@@ -0,0 +1 @@
+

trustgraph_span-2.11.1/trustgraph/model/__init__.py

@@ -0,0 +1 @@
+

trustgraph_span-2.11.1/trustgraph/model/field_extraction/__init__.py

@@ -0,0 +1 @@
+

trustgraph_span-2.11.1/trustgraph/model/field_extraction/span/__init__.py

@@ -0,0 +1,2 @@
+
+from . service import *

trustgraph_span-2.11.1/trustgraph/model/field_extraction/span/__main__.py

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

trustgraph_span-2.11.1/trustgraph/model/field_extraction/span/engine.py

@@ -0,0 +1,157 @@
+
+"""
+ONNX inference engine for BIO span extraction.
+Loads a span-extraction model from HuggingFace Hub and runs
+instruction-based field extraction using BIO tagging.
+"""
+
+import json
+import logging
+from typing import Any, Dict, List, Optional, Union
+
+import numpy as np
+import onnxruntime as ort
+from tokenizers import Tokenizer
+from huggingface_hub import hf_hub_download
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_REPO = "tasksource/UltraBERT-150M-clairic-extraction"
+DEFAULT_MAX_LENGTH = 512
+
+O, B, I = 0, 1, 2
+
+
+def _serialize_state(state):
+    if isinstance(state, str):
+        return state
+    return json.dumps(state, ensure_ascii=False)
+
+
+def _decode_spans(text, tags, offsets, seq_ids):
+    """Extract text spans from BIO tags.
+
+    Only tokens in the text sequence (seq_id=1) are considered.
+    """
+    spans = []
+    start = None
+
+    for i, tag in enumerate(tags):
+        in_text = (seq_ids[i] == 1)
+
+        if start is not None and (not in_text or tag != I):
+            char_start = offsets[start][0]
+            char_end = offsets[i - 1][1]
+            span = text[char_start:char_end].strip().rstrip(",;:")
+            if span:
+                spans.append(span)
+            start = None
+
+        if in_text and tag in (B, I) and start is None:
+            start = i
+
+    if start is not None:
+        char_start = offsets[start][0]
+        last = len(tags) - 1
+        while last > start and seq_ids[last] != 1:
+            last -= 1
+        char_end = offsets[last][1]
+        span = text[char_start:char_end].strip().rstrip(",;:")
+        if span:
+            spans.append(span)
+
+    return spans
+
+
+class OnnxEngine:
+
+    def __init__(
+        self,
+        repo_id: str = DEFAULT_REPO,
+        local_dir: Optional[str] = None,
+        max_length: int = DEFAULT_MAX_LENGTH,
+    ):
+        local_dir = local_dir or "/opt/models/span-extraction"
+
+        logger.info(
+            "Downloading span-extraction model from %s ...", repo_id,
+        )
+        model_path = hf_hub_download(
+            repo_id=repo_id, filename="onnx/model_quantized.onnx",
+            local_dir=local_dir,
+        )
+        tokenizer_path = hf_hub_download(
+            repo_id=repo_id, filename="tokenizer.json",
+            local_dir=local_dir,
+        )
+
+        logger.info("Loading ONNX session...")
+        self.session = ort.InferenceSession(
+            model_path, providers=["CPUExecutionProvider"],
+        )
+        self.tokenizer = Tokenizer.from_file(tokenizer_path)
+        self.tokenizer.no_padding()
+
+        self.max_length = max_length
+        self.repo_id = repo_id
+
+        logger.info(
+            "Span-extraction engine loaded (max_length=%d)", max_length,
+        )
+
+    def predict(
+        self,
+        state: Union[str, dict, list],
+        fields: Dict[str, str],
+    ) -> Dict[str, Any]:
+
+        text = _serialize_state(state)
+
+        field_names = list(fields.keys())
+        instructions = list(fields.values())
+
+        encodings = [
+            self.tokenizer.encode(inst, text) for inst in instructions
+        ]
+
+        if self.max_length:
+            for enc in encodings:
+                enc.truncate(self.max_length)
+
+        max_len = max(len(enc.ids) for enc in encodings)
+        pad_id = 0
+
+        all_input_ids = []
+        all_attention_mask = []
+
+        for enc in encodings:
+            ids = list(enc.ids)
+            pad_len = max_len - len(ids)
+            all_input_ids.append(ids + [pad_id] * pad_len)
+            all_attention_mask.append([1] * len(ids) + [0] * pad_len)
+
+        input_ids = np.array(all_input_ids, dtype=np.int64)
+        attention_mask = np.array(all_attention_mask, dtype=np.int64)
+
+        outputs = self.session.run(None, {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+        })
+
+        tags = outputs[0].argmax(axis=-1)
+
+        answers = {}
+        for i, name in enumerate(field_names):
+            enc = encodings[i]
+            n_tokens = len(enc.ids)
+            answers[name] = _decode_spans(
+                text,
+                tags[i][:n_tokens].tolist(),
+                enc.offsets,
+                enc.sequence_ids,
+            )
+
+        return {
+            "model": self.repo_id,
+            "answers": answers,
+        }

trustgraph_span-2.11.1/trustgraph/model/field_extraction/span/service.py

@@ -0,0 +1,136 @@
+
+"""
+Field extraction service using BIO span extraction via ONNX Runtime.
+Extracts verbatim text spans from input based on field instructions.
+"""
+
+import asyncio
+import logging
+
+from .... base import FlowProcessor, ConsumerSpec, ProducerSpec, ParameterSpec
+from .... schema import FieldExtractionRequest, FieldExtractionResponse
+from .... schema import Error
+
+from . engine import OnnxEngine, DEFAULT_REPO, DEFAULT_MAX_LENGTH
+
+logger = logging.getLogger(__name__)
+
+default_ident = "field-extraction"
+
+
+class Processor(FlowProcessor):
+
+    def __init__(self, **params):
+
+        super(Processor, self).__init__(**params)
+
+        self.repo_id = params.get("model") or DEFAULT_REPO
+        self.max_length = int(
+            params.get("max-length") or DEFAULT_MAX_LENGTH
+        )
+
+        self.register_specification(
+            ParameterSpec(
+                name="model",
+            )
+        )
+
+        self.register_specification(
+            ParameterSpec(
+                name="max-length",
+            )
+        )
+
+        self.register_specification(
+            ConsumerSpec(
+                name="request",
+                schema=FieldExtractionRequest,
+                handler=self.on_request,
+            )
+        )
+
+        self.register_specification(
+            ProducerSpec(
+                name="response",
+                schema=FieldExtractionResponse,
+            )
+        )
+
+        self.engine = None
+
+    async def start(self):
+
+        logger.info("Loading span-extraction ONNX engine...")
+        loop = asyncio.get_event_loop()
+        self.engine = await loop.run_in_executor(
+            None,
+            lambda: OnnxEngine(
+                repo_id=self.repo_id, max_length=self.max_length,
+            ),
+        )
+        logger.info("Span-extraction ONNX engine loaded successfully")
+
+        await super().start()
+
+    async def on_request(self, msg, consumer, flow):
+
+        request = msg.value()
+        id = msg.properties()["id"]
+        producer = flow("response")
+
+        if self.engine is None:
+            response = FieldExtractionResponse(
+                error=Error(
+                    type="not-ready",
+                    message="Span-extraction engine is still loading",
+                ),
+            )
+            await producer.send(response, properties={"id": id})
+            return
+
+        try:
+
+            result = self.engine.predict(
+                request.state, request.fields,
+            )
+
+            response = FieldExtractionResponse(
+                answers=result.get("answers", {}),
+                model=result.get("model", ""),
+            )
+
+        except Exception as e:
+            logger.error(f"Field extraction error: {e}")
+            response = FieldExtractionResponse(
+                error=Error(
+                    type="field-extraction-error",
+                    message=str(e),
+                ),
+            )
+
+        await producer.send(response, properties={"id": id})
+
+    @staticmethod
+    def add_args(parser):
+
+        FlowProcessor.add_args(parser)
+
+        parser.add_argument(
+            '--model',
+            default=None,
+            help='HuggingFace repo for ONNX model (default: %s)'
+            % DEFAULT_REPO,
+        )
+
+        parser.add_argument(
+            '--max-length',
+            default=None,
+            type=int,
+            help='Max token sequence length (default: %d)'
+            % DEFAULT_MAX_LENGTH,
+        )
+
+
+def run():
+
+    Processor.launch(default_ident, __doc__)

trustgraph_span-2.11.1/trustgraph/span_version.py

@@ -0,0 +1 @@
+__version__ = "2.11.1"

trustgraph_span-2.11.1/trustgraph_span.egg-info/PKG-INFO

@@ -0,0 +1,19 @@
+Metadata-Version: 2.4
+Name: trustgraph-span
+Version: 2.11.1
+Summary: Span-extraction field extraction service for TrustGraph.
+Author-email: "trustgraph.ai" <security@trustgraph.ai>
+Project-URL: Homepage, https://github.com/trustgraph-ai/trustgraph
+Classifier: Programming Language :: Python :: 3
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: trustgraph-base<2.12,>=2.11
+Requires-Dist: trustgraph-flow<2.12,>=2.11
+Requires-Dist: onnxruntime
+Requires-Dist: tokenizers
+Requires-Dist: huggingface-hub
+Requires-Dist: numpy
+Requires-Dist: pulsar-client
+Requires-Dist: pyyaml
+Requires-Dist: prometheus-client

trustgraph_span-2.11.1/trustgraph_span.egg-info/SOURCES.txt

@@ -0,0 +1,15 @@
+pyproject.toml
+trustgraph/__init__.py
+trustgraph/span_version.py
+trustgraph/model/__init__.py
+trustgraph/model/field_extraction/__init__.py
+trustgraph/model/field_extraction/span/__init__.py
+trustgraph/model/field_extraction/span/__main__.py
+trustgraph/model/field_extraction/span/engine.py
+trustgraph/model/field_extraction/span/service.py
+trustgraph_span.egg-info/PKG-INFO
+trustgraph_span.egg-info/SOURCES.txt
+trustgraph_span.egg-info/dependency_links.txt
+trustgraph_span.egg-info/entry_points.txt
+trustgraph_span.egg-info/requires.txt
+trustgraph_span.egg-info/top_level.txt

trustgraph_span-2.11.1/trustgraph_span.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

trustgraph_span-2.11.1/trustgraph_span.egg-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+field-extraction-span = trustgraph.model.field_extraction.span:run

trustgraph_span-2.11.1/trustgraph_span.egg-info/requires.txt

@@ -0,0 +1,9 @@
+trustgraph-base<2.12,>=2.11
+trustgraph-flow<2.12,>=2.11
+onnxruntime
+tokenizers
+huggingface-hub
+numpy
+pulsar-client
+pyyaml
+prometheus-client

trustgraph_span-2.11.1/trustgraph_span.egg-info/top_level.txt

@@ -0,0 +1 @@
+trustgraph