glinker 0.1.0__py3-none-any.whl

glinker/l3/processor.py

@@ -0,0 +1,350 @@
+from typing import Any, List, Optional
+import torch
+from glinker.core.base import BaseProcessor
+from glinker.core.registry import processor_registry
+from .models import L3Config, L3Input, L3Output, L3Entity
+from .component import L3Component
+
+
+class L3Processor(BaseProcessor[L3Config, L3Input, L3Output]):
+    """GLiNER entity linking processor"""
+
+    def __init__(
+        self,
+        config: L3Config,
+        component: L3Component,
+        pipeline: list[tuple[str, dict[str, Any]]] = None
+    ):
+        super().__init__(config, component, pipeline)
+        self._validate_pipeline()
+        self.schema = {}
+        self._l2_processor = None  # Will be set by DAG executor for cache write-back
+    
+    def _default_pipeline(self) -> list[tuple[str, dict[str, Any]]]:
+        return [
+            ("predict_entities", {}),
+            ("filter_by_score", {}),
+            ("sort_by_position", {})
+        ]
+    
+    @staticmethod
+    def _build_input_spans(l1_entities_for_text: List[Any]) -> List[List[dict]]:
+        """Convert L1 entities to GLiNER input_spans format.
+
+        Args:
+            l1_entities_for_text: List of L1Entity objects for a single text
+
+        Returns:
+            List of span dicts with 'start' and 'end' keys,
+            wrapped in an outer list as expected by GLiNER input_spans.
+        """
+        spans = [{"start": e.start, "end": e.end} for e in l1_entities_for_text]
+        return [spans]
+
+    def __call__(
+        self,
+        texts: List[str] = None,
+        candidates: List[List[Any]] = None,
+        l1_entities: List[List[Any]] = None,
+        input_data: L3Input = None
+    ) -> L3Output:
+        """Process texts with candidate labels
+
+        Args:
+            texts: List of input texts
+            candidates: List of candidate lists per text (from L2)
+            l1_entities: Optional L1 entities per text, used to build input_spans
+                         so L3 predicts on the same spans extracted in L1
+            input_data: Alternative L3Input object
+        """
+
+        # Support both direct params and L3Input
+        if texts is not None and candidates is not None:
+            texts_to_process = texts
+            candidates_to_process = candidates
+        elif input_data is not None:
+            texts_to_process = input_data.texts
+            candidates_to_process = input_data.labels
+        else:
+            raise ValueError("Either 'texts'+'candidates' or 'input_data' must be provided")
+
+        all_entities = []
+
+        # Detect shared candidates (all texts use the same list, e.g. simple pipeline)
+        shared = (
+            len(candidates_to_process) > 1
+            and all(c is candidates_to_process[0] for c in candidates_to_process[1:])
+        )
+
+        # Pre-compute labels & embeddings once when candidates are shared
+        shared_labels = None
+        shared_label_to_candidate = None
+        shared_use_precomputed = False
+        shared_embeddings = None
+
+        if shared:
+            ref_candidates = candidates_to_process[0]
+            if self.schema:
+                shared_labels, shared_label_to_candidate = (
+                    self._create_gliner_labels_with_mapping(ref_candidates)
+                )
+            else:
+                shared_labels = [self._extract_label(c) for c in ref_candidates]
+                shared_label_to_candidate = {}
+
+            shared_use_precomputed = (
+                self.config.use_precomputed_embeddings
+                and self.component.supports_precomputed_embeddings
+                and self._can_use_precomputed(ref_candidates, shared_label_to_candidate)
+            )
+            if shared_use_precomputed:
+                shared_embeddings = self._get_embeddings_tensor(
+                    ref_candidates, shared_labels, shared_label_to_candidate
+                )
+
+        for idx, (text, text_candidates) in enumerate(zip(texts_to_process, candidates_to_process)):
+            # Build input_spans from L1 entities if available
+            input_spans = None
+            if l1_entities is not None and idx < len(l1_entities):
+                text_l1 = l1_entities[idx]
+                if text_l1:
+                    input_spans = self._build_input_spans(text_l1)
+
+            if shared:
+                labels = shared_labels
+                label_to_candidate = shared_label_to_candidate
+                use_precomputed = shared_use_precomputed
+                embeddings = shared_embeddings
+            else:
+                # Create labels from candidates (per-text)
+                if self.schema:
+                    labels, label_to_candidate = self._create_gliner_labels_with_mapping(text_candidates)
+                else:
+                    labels = [self._extract_label(c) for c in text_candidates]
+                    label_to_candidate = {}
+
+                use_precomputed = (
+                    self.config.use_precomputed_embeddings
+                    and self.component.supports_precomputed_embeddings
+                    and self._can_use_precomputed(text_candidates, label_to_candidate)
+                )
+                embeddings = None
+                if use_precomputed:
+                    embeddings = self._get_embeddings_tensor(text_candidates, labels, label_to_candidate)
+
+            if use_precomputed and embeddings is not None:
+                entities = self.component.predict_with_embeddings(
+                    text, labels, embeddings, input_spans=input_spans
+                )
+            else:
+                # Regular prediction
+                entities = self.component.predict_entities(text, labels, input_spans=input_spans)
+
+                # Optionally cache computed embeddings
+                if self.config.cache_embeddings and self.component.supports_precomputed_embeddings:
+                    self._cache_embeddings(text_candidates, labels, label_to_candidate)
+
+            # Apply rest of pipeline
+            for method_name, kwargs in self.pipeline[1:]:
+                method = getattr(self.component, method_name)
+                entities = method(entities, **kwargs)
+
+            # Apply ranking if configured
+            if self.schema.get('ranking'):
+                entities = self._rank_entities(entities, text_candidates)
+
+            all_entities.append(entities)
+
+        return L3Output(entities=all_entities)
+
+    def _can_use_precomputed(
+        self,
+        candidates: List[Any],
+        label_to_candidate: dict
+    ) -> bool:
+        """Check if all candidates have compatible precomputed embeddings"""
+        if not candidates:
+            return False
+
+        expected_model = self.config.model_name
+
+        for candidate in candidates:
+            # Check if candidate has embedding
+            embedding = getattr(candidate, 'embedding', None)
+            if embedding is None:
+                return False
+
+            # Check if model matches
+            model_id = getattr(candidate, 'embedding_model_id', None)
+            if model_id != expected_model:
+                return False
+
+        return True
+
+    def _get_embeddings_tensor(
+        self,
+        candidates: List[Any],
+        labels: List[str],
+        label_to_candidate: dict
+    ) -> torch.Tensor:
+        """Build embeddings tensor from candidates in same order as labels"""
+        embeddings = []
+
+        for label in labels:
+            candidate = label_to_candidate.get(label)
+            if candidate and hasattr(candidate, 'embedding') and candidate.embedding:
+                embeddings.append(candidate.embedding)
+            else:
+                # Should not happen if _can_use_precomputed returned True
+                raise ValueError(f"Missing embedding for label: {label}")
+
+        return torch.tensor(embeddings, device=self.component.device)
+
+    def _cache_embeddings(
+        self,
+        candidates: List[Any],
+        labels: List[str],
+        label_to_candidate: dict
+    ):
+        """Compute and cache embeddings for candidates without them"""
+        if not self._l2_processor:
+            return
+
+        # Find candidates without embeddings
+        to_compute = []
+        to_compute_ids = []
+
+        for candidate in candidates:
+            if not getattr(candidate, 'embedding', None):
+                to_compute.append(candidate)
+                to_compute_ids.append(candidate.entity_id)
+
+        if not to_compute:
+            return
+
+        # Format labels for these candidates
+        template = self.schema.get('template', '{label}')
+        compute_labels = []
+        for candidate in to_compute:
+            try:
+                if hasattr(candidate, 'model_dump'):
+                    formatted = template.format(**candidate.model_dump())
+                elif hasattr(candidate, 'dict'):
+                    formatted = template.format(**candidate.dict())
+                else:
+                    formatted = candidate.label
+                compute_labels.append(formatted)
+            except KeyError:
+                compute_labels.append(candidate.label)
+
+        # Encode labels
+        embeddings = self.component.encode_labels(compute_labels)
+
+        # Update L2 layer
+        if hasattr(self._l2_processor, 'component'):
+            for layer in self._l2_processor.component.layers:
+                if layer.is_available():
+                    layer.update_embeddings(
+                        to_compute_ids,
+                        embeddings.tolist(),
+                        self.config.model_name
+                    )
+                    break  # Update first available layer
+    
+    def _extract_label(self, candidate: Any) -> str:
+        """Extract label from candidate"""
+        if hasattr(candidate, 'label'):
+            return candidate.label
+        return str(candidate)
+
+    def _create_gliner_labels_with_mapping(self, candidates: List[Any]) -> tuple:
+        """
+        Create GLiNER labels using schema template and return label->candidate mapping.
+
+        Returns:
+            tuple: (labels: List[str], label_to_candidate: dict)
+        """
+        template = self.schema.get('template', '{label}')
+        labels = []
+        label_to_candidate = {}
+        seen = set()
+
+        for candidate in candidates:
+            try:
+                if hasattr(candidate, 'model_dump'):
+                    cand_dict = candidate.model_dump()
+                elif hasattr(candidate, 'dict'):
+                    cand_dict = candidate.dict()
+                elif isinstance(candidate, dict):
+                    cand_dict = candidate
+                else:
+                    label = str(candidate)
+                    if label.lower() not in seen:
+                        labels.append(label)
+                        seen.add(label.lower())
+                    continue
+
+                label = template.format(**cand_dict)
+                label_lower = label.lower()
+                if label_lower not in seen:
+                    labels.append(label)
+                    label_to_candidate[label] = candidate
+                    seen.add(label_lower)
+            except (KeyError, AttributeError):
+                if hasattr(candidate, 'label'):
+                    if candidate.label.lower() not in seen:
+                        labels.append(candidate.label)
+                        label_to_candidate[candidate.label] = candidate
+                        seen.add(candidate.label.lower())
+
+        return labels, label_to_candidate
+
+    def _create_gliner_labels(self, candidates: List[Any]) -> List[str]:
+        """Create GLiNER labels using schema template (legacy, for compatibility)"""
+        labels, _ = self._create_gliner_labels_with_mapping(candidates)
+        return labels
+    
+    def _rank_entities(self, entities: List[L3Entity], candidates: List[Any]) -> List[L3Entity]:
+        """Re-rank entities using multiple scoring factors"""
+        # Build label to candidate mapping
+        label_to_candidate = {}
+        for c in candidates:
+            if hasattr(c, 'label'):
+                label_to_candidate[c.label] = c
+                if hasattr(c, 'aliases'):
+                    for alias in c.aliases:
+                        if alias not in label_to_candidate:
+                            label_to_candidate[alias] = c
+        
+        # Calculate weighted scores
+        for entity in entities:
+            total_score = 0.0
+            total_weight = 0.0
+            
+            for rank_spec in self.schema['ranking']:
+                field = rank_spec['field']
+                weight = rank_spec['weight']
+                total_weight += weight
+                
+                if field == 'gliner_score':
+                    total_score += entity.score * weight
+                else:
+                    candidate = label_to_candidate.get(entity.label)
+                    if candidate and hasattr(candidate, field):
+                        value = getattr(candidate, field, 0)
+                        if isinstance(value, (int, float)):
+                            normalized = min(value / 1000000.0, 1.0)
+                            total_score += normalized * weight
+            
+            if total_weight > 0:
+                entity.score = total_score / total_weight
+        
+        return sorted(entities, key=lambda x: x.score, reverse=True)
+
+
+@processor_registry.register("l3_batch")
+def create_l3_processor(config_dict: dict, pipeline: list = None) -> L3Processor:
+    """Factory: creates component + processor"""
+    config = L3Config(**config_dict)
+    component = L3Component(config)
+    return L3Processor(config, component, pipeline)

glinker/l4/__init__.py

@@ -0,0 +1,9 @@
+from .models import L4Config
+from .component import L4Component
+from .processor import L4Processor
+
+__all__ = [
+    "L4Config",
+    "L4Component",
+    "L4Processor",
+]

glinker/l4/component.py

@@ -0,0 +1,121 @@
+from typing import List, Optional
+from gliner import GLiNER
+from glinker.core.base import BaseComponent
+from glinker.l3.models import L3Entity
+from .models import L4Config
+
+
+class L4Component(BaseComponent[L4Config]):
+    """GLiNER-based reranking component (uni-encoder only, no precomputed embeddings)"""
+
+    def _setup(self):
+        """Initialize GLiNER model"""
+        self.model = GLiNER.from_pretrained(
+            self.config.model_name,
+            token=self.config.token,
+            max_length=self.config.max_length
+        )
+        self.model.to(self.config.device)
+
+    def get_available_methods(self) -> List[str]:
+        return [
+            "predict_entities",
+            "predict_entities_chunked",
+            "filter_by_score",
+            "sort_by_position",
+            "deduplicate_entities"
+        ]
+
+    def predict_entities(
+        self,
+        text: str,
+        labels: List[str],
+        input_spans: List[List[dict]] = None
+    ) -> List[L3Entity]:
+        """Predict entities using GLiNER for a single label set.
+
+        Args:
+            text: Input text
+            labels: List of label strings
+            input_spans: Optional list of span dicts with 'start' and 'end' keys
+        """
+        if not labels:
+            return []
+
+        kwargs = dict(
+            threshold=self.config.threshold,
+            flat_ner=self.config.flat_ner,
+            multi_label=self.config.multi_label,
+            return_class_probs=True
+        )
+        if input_spans is not None:
+            kwargs["input_spans"] = input_spans
+
+        entities = self.model.predict_entities(text, labels, **kwargs)
+
+        return [
+            L3Entity(
+                text=e["text"],
+                label=e["label"],
+                start=e["start"],
+                end=e["end"],
+                score=e["score"],
+                class_probs=e.get("class_probs")
+            )
+            for e in entities
+        ]
+
+    def predict_entities_chunked(
+        self,
+        text: str,
+        labels: List[str],
+        max_labels: int,
+        input_spans: List[List[dict]] = None
+    ) -> List[L3Entity]:
+        """Predict entities with candidate chunking.
+
+        Splits labels into chunks of max_labels, runs inference on each chunk,
+        and merges results.
+
+        Args:
+            text: Input text
+            labels: Full list of candidate label strings
+            max_labels: Maximum labels per inference call
+            input_spans: Optional span constraints from L1 entities
+        """
+        if not labels:
+            return []
+
+        if len(labels) <= max_labels:
+            return self.predict_entities(text, labels, input_spans=input_spans)
+
+        # Split labels into chunks
+        chunks = [
+            labels[i:i + max_labels]
+            for i in range(0, len(labels), max_labels)
+        ]
+
+        all_entities = []
+        for chunk in chunks:
+            entities = self.predict_entities(text, chunk, input_spans=input_spans)
+            all_entities.extend(entities)
+
+        return all_entities
+
+    def filter_by_score(self, entities: List[L3Entity], threshold: float = None) -> List[L3Entity]:
+        """Filter entities by confidence score"""
+        threshold = threshold if threshold is not None else self.config.threshold
+        return [e for e in entities if e.score >= threshold]
+
+    def sort_by_position(self, entities: List[L3Entity]) -> List[L3Entity]:
+        """Sort entities by position in text"""
+        return sorted(entities, key=lambda e: e.start)
+
+    def deduplicate_entities(self, entities: List[L3Entity]) -> List[L3Entity]:
+        """Remove duplicate entities, keeping the highest-scoring one per span"""
+        best = {}
+        for entity in entities:
+            key = (entity.text, entity.start, entity.end)
+            if key not in best or entity.score > best[key].score:
+                best[key] = entity
+        return list(best.values())

glinker/l4/models.py

@@ -0,0 +1,21 @@
+from pydantic import Field
+from typing import List, Any, Optional
+from glinker.core.base import BaseConfig, BaseInput, BaseOutput
+
+
+class L4Config(BaseConfig):
+    model_name: str = Field(...)
+    token: str = Field(None)
+    device: str = Field("cpu")
+    threshold: float = Field(0.5)
+    flat_ner: bool = Field(True)
+    multi_label: bool = Field(False)
+    max_labels: int = Field(
+        20,
+        description="Maximum number of candidate labels per inference call. "
+        "When candidates exceed this, they are split into chunks."
+    )
+    max_length: int = Field(
+        None,
+        description="Maximum sequence length for tokenization. Passed to GLiNER.from_pretrained."
+    )

glinker/l4/processor.py

@@ -0,0 +1,156 @@
+from typing import Any, List, Optional
+from glinker.core.base import BaseProcessor
+from glinker.core.registry import processor_registry
+from glinker.l3.models import L3Input, L3Output, L3Entity
+from .models import L4Config
+from .component import L4Component
+
+
+class L4Processor(BaseProcessor[L4Config, L3Input, L3Output]):
+    """GLiNER reranking processor with candidate chunking"""
+
+    def __init__(
+        self,
+        config: L4Config,
+        component: L4Component,
+        pipeline: list[tuple[str, dict[str, Any]]] = None
+    ):
+        super().__init__(config, component, pipeline)
+        self._validate_pipeline()
+        self.schema = {}
+
+    def _default_pipeline(self) -> list[tuple[str, dict[str, Any]]]:
+        return [
+            ("predict_entities_chunked", {}),
+            ("deduplicate_entities", {}),
+            ("filter_by_score", {}),
+            ("sort_by_position", {})
+        ]
+
+    @staticmethod
+    def _build_input_spans(l1_entities_for_text: List[Any]) -> List[List[dict]]:
+        """Convert L1 entities to GLiNER input_spans format."""
+        spans = [{"start": e.start, "end": e.end} for e in l1_entities_for_text]
+        return [spans]
+
+    def __call__(
+        self,
+        texts: List[str] = None,
+        candidates: List[List[Any]] = None,
+        l1_entities: List[List[Any]] = None,
+        input_data: L3Input = None
+    ) -> L3Output:
+        """Process texts with candidate labels using chunked GLiNER inference.
+
+        Args:
+            texts: List of input texts
+            candidates: List of candidate lists per text (from L2)
+            l1_entities: Optional L1 entities per text, used to build input_spans
+            input_data: Alternative L3Input object
+        """
+        if texts is not None and candidates is not None:
+            texts_to_process = texts
+            candidates_to_process = candidates
+        elif input_data is not None:
+            texts_to_process = input_data.texts
+            candidates_to_process = input_data.labels
+        else:
+            raise ValueError("Either 'texts'+'candidates' or 'input_data' must be provided")
+
+        all_entities = []
+        max_labels = self.config.max_labels
+
+        # Detect shared candidates (all texts use the same list)
+        shared = (
+            len(candidates_to_process) > 1
+            and all(c is candidates_to_process[0] for c in candidates_to_process[1:])
+        )
+
+        shared_labels = None
+        if shared:
+            ref_candidates = candidates_to_process[0]
+            if self.schema:
+                shared_labels, _ = self._create_gliner_labels_with_mapping(ref_candidates)
+            else:
+                shared_labels = [self._extract_label(c) for c in ref_candidates]
+
+        for idx, (text, text_candidates) in enumerate(zip(texts_to_process, candidates_to_process)):
+            # Build input_spans from L1 entities if available
+            input_spans = None
+            if l1_entities is not None and idx < len(l1_entities):
+                text_l1 = l1_entities[idx]
+                if text_l1:
+                    input_spans = self._build_input_spans(text_l1)
+
+            if shared:
+                labels = shared_labels
+            else:
+                if self.schema:
+                    labels, _ = self._create_gliner_labels_with_mapping(text_candidates)
+                else:
+                    labels = [self._extract_label(c) for c in text_candidates]
+
+            # Run chunked prediction
+            entities = self.component.predict_entities_chunked(
+                text, labels, max_labels, input_spans=input_spans
+            )
+
+            # Apply remaining pipeline steps (deduplicate, filter, sort)
+            for method_name, kwargs in self.pipeline[1:]:
+                method = getattr(self.component, method_name)
+                entities = method(entities, **kwargs)
+
+            all_entities.append(entities)
+
+        return L3Output(entities=all_entities)
+
+    def _extract_label(self, candidate: Any) -> str:
+        """Extract label from candidate"""
+        if hasattr(candidate, 'label'):
+            return candidate.label
+        return str(candidate)
+
+    def _create_gliner_labels_with_mapping(self, candidates: List[Any]) -> tuple:
+        """Create GLiNER labels using schema template and return label->candidate mapping."""
+        template = self.schema.get('template', '{label}')
+        labels = []
+        label_to_candidate = {}
+        seen = set()
+
+        for candidate in candidates:
+            try:
+                if hasattr(candidate, 'model_dump'):
+                    cand_dict = candidate.model_dump()
+                elif hasattr(candidate, 'dict'):
+                    cand_dict = candidate.dict()
+                elif isinstance(candidate, dict):
+                    cand_dict = candidate
+                else:
+                    label = str(candidate)
+                    if label.lower() not in seen:
+                        labels.append(label)
+                        seen.add(label.lower())
+                    continue
+
+                label = template.format(**cand_dict)
+                label_lower = label.lower()
+                if label_lower not in seen:
+                    labels.append(label)
+                    label_to_candidate[label] = candidate
+                    seen.add(label_lower)
+            except (KeyError, AttributeError):
+                if hasattr(candidate, 'label'):
+                    if candidate.label.lower() not in seen:
+                        labels.append(candidate.label)
+                        label_to_candidate[candidate.label] = candidate
+                        seen.add(candidate.label.lower())
+
+        return labels, label_to_candidate
+
+
+@processor_registry.register("l4_reranker")
+def create_l4_processor(config_dict: dict, pipeline: list = None) -> L4Processor:
+    """Factory: creates component + processor"""
+    config = L4Config(**config_dict)
+    component = L4Component(config)
+    return L4Processor(config, component, pipeline)

glinker/py.typed

@@ -0,0 +1 @@
+# Marker file for PEP 561