diff --git a/README.md b/README.md index 96a460a..1b9428f 100644 --- a/README.md +++ b/README.md @@ -129,27 +129,25 @@ builder.l3.configure(model="knowledgator/gliner-linker-large-v1.0") executor = DAGExecutor(builder.get_config()) # 3. Load entities -executor.load_entities("data/entities.jsonl", target_layers=["dict"]) +executor.load_entities("data/pubmesh_ontology.jsonl", target_layers=["dict"]) # 4. Process text result = executor.execute({ - "texts": ["Farnese Palace is one of the most important palaces in the city of Rome."] + "texts": ["BRCA1 mutations are associated with breast cancer risk."] }) # 5. Get results l0_result = result.get("l0_result") -for entity in l0_result.entities: - if entity.linked_entity: - print(f"{entity.mention_text} → {entity.linked_entity.label}") - print(f" Confidence: {entity.linked_entity.score:.3f}") +for text_entities in l0_result.entities: + for entity in text_entities: + if entity.linked_entity: + print(f"{entity.mention_text} → {entity.linked_entity.label}") + print(f" Confidence: {entity.linked_entity.confidence:.3f}") ``` **Output:** ``` -BRCA1 → BRCA1: Breast cancer type 1 susceptibility protein - Confidence: 0.923 -breast cancer → Breast Cancer: Malignant neoplasm of the breast - Confidence: 0.887 +# Labels and confidence values depend on the model and loaded entity database. ``` --- @@ -284,7 +282,7 @@ For full control over every layer, define the pipeline in YAML and load it: ```python from glinker import ProcessorFactory -executor = ProcessorFactory.create_pipeline("configs/pipelines/dict/simple.yaml") +executor = ProcessorFactory.create_pipeline("configs/pipelines/dict/default.yaml") executor.load_entities("data/entities.jsonl") result = executor.execute({"texts": ["TP53 mutations cause cancer"]}) ``` diff --git a/configs/pipelines/README.md b/configs/pipelines/README.md index cb9558f..9e92d43 100644 --- a/configs/pipelines/README.md +++ b/configs/pipelines/README.md @@ -59,13 +59,11 @@ python scripts/database/jsonl2postgresql.py - Use case: High-performance BiEncoder inference ```python -from src.core.dag import DAGExecutor, DAGPipeline -import yaml +from glinker import ProcessorFactory -with open("configs/pipelines/postgres/with_embeddings.yaml") as f: - pipeline = DAGPipeline(**yaml.safe_load(f)) - -executor = DAGExecutor(pipeline) +executor = ProcessorFactory.create_pipeline( + "configs/pipelines/postgres/with_embeddings.yaml" +) # Step 1: Load entities executor.load_entities("entities.jsonl", target_layers=["postgres"]) @@ -233,13 +231,11 @@ python demo.py -c configs/pipelines/postgres_redis/default.yaml ```bash # Setup embeddings python -c " -from src.core.dag import DAGExecutor, DAGPipeline -import yaml +from glinker import ProcessorFactory -with open('configs/pipelines/postgres/with_embeddings.yaml') as f: - pipeline = DAGPipeline(**yaml.safe_load(f)) - -executor = DAGExecutor(pipeline) +executor = ProcessorFactory.create_pipeline( + 'configs/pipelines/postgres/with_embeddings.yaml' +) executor.load_entities('entities.jsonl', target_layers=['postgres']) executor.precompute_embeddings(target_layers=['postgres']) " diff --git a/docs/DATABASE_SETUP.md b/docs/DATABASE_SETUP.md index 0834af7..080701e 100644 --- a/docs/DATABASE_SETUP.md +++ b/docs/DATABASE_SETUP.md @@ -351,7 +351,7 @@ from glinker.core.factory import ProcessorFactory import yaml # Load pipeline config with multiple database layers -with open("configs/pipelines/postgres_redis_basic.yaml") as f: +with open("configs/pipelines/postgres_redis/default.yaml") as f: config = yaml.safe_load(f) executor = ProcessorFactory.create_from_dict(config, verbose=True) @@ -379,7 +379,7 @@ from glinker.core.factory import ProcessorFactory import yaml # Use embedding-enabled config -with open("configs/pipelines/postgres_redis_embeddings.yaml") as f: +with open("configs/pipelines/postgres_redis/with_embeddings.yaml") as f: config = yaml.safe_load(f) executor = ProcessorFactory.create_from_dict(config, verbose=True) diff --git a/src/glinker/core/__init__.py b/src/glinker/core/__init__.py index cee531c..954336e 100644 --- a/src/glinker/core/__init__.py +++ b/src/glinker/core/__init__.py @@ -21,6 +21,7 @@ from .factory import ProcessorFactory, load_yaml from .builders import ConfigBuilder from .registry import ProcessorRegistry, processor_registry +from .gliner_model import GLinkerModel, SparseTokenModel __all__ = [ "BaseComponent", @@ -33,6 +34,7 @@ "DAGExecutor", "DAGPipeline", "FieldResolver", + "GLinkerModel", "InputConfig", "InputT", "OutputConfig", @@ -42,6 +44,7 @@ "ProcessorFactory", "ProcessorRegistry", "ReshapeConfig", + "SparseTokenModel", "load_yaml", "processor_registry", ] diff --git a/src/glinker/core/gliner_model.py b/src/glinker/core/gliner_model.py new file mode 100644 index 0000000..9e890ff --- /dev/null +++ b/src/glinker/core/gliner_model.py @@ -0,0 +1,578 @@ +"""GLinkerModel — BiEncoderTokenGLiNER subclass with sparse per-span scoring. + +Architecture +------------ +SparseTokenModel subclasses BiEncoderTokenModel and overrides forward(). +The scorer itself is untouched — sparse packing/unpacking wraps around it: + + preprocess_sparse(token_rep, label_rep) + → packs M mentions x their candidates into (M_total, max_sl, max_lc, H) + + self.scorer(super_tok, super_lbl) + → (M_total, max_sl, max_lc, 3) ← one GPU call vs O(seq x labels) + + postprocess_sparse(packed_scores, meta) + → full (B, seq_len, num_classes, 3), fill=-10 for unscored positions + +Without sparse config (_word_spans is None), forward is identical to parent. +""" + +from __future__ import annotations + +import logging +from typing import Any, Dict, List, Tuple + +import torch +from gliner.modeling.base import BiEncoderTokenModel +from gliner.modeling.outputs import GLiNERBaseOutput + +try: + from gliner.model import BiEncoderTokenGLiNER +except ImportError: + BiEncoderTokenGLiNER = None + +logger = logging.getLogger(__name__) + +SPARSE_FILL: float = -10.0 # sigmoid(-10) ≈ 0 — below any practical threshold + + +# --------------------------------------------------------------------------- +# SparseTokenModel +# --------------------------------------------------------------------------- + + +def _pad_dim(t: torch.Tensor, target: int, dim: int) -> torch.Tensor: + """Pad tensor along *dim* to *target* length with zeros.""" + delta = target - t.shape[dim] + if delta <= 0: + return t + pad_shape = list(t.shape) + pad_shape[dim] = delta + return torch.cat([t, t.new_zeros(pad_shape)], dim=dim) + + +class SparseTokenModel(BiEncoderTokenModel): + """BiEncoderTokenModel with built-in sparse per-span scoring. + + Two new methods handle the optimisation: + + ``preprocess_sparse(token_rep, label_rep)`` + For each batch item, gathers only the word positions that belong to a + candidate mention and only the label embeddings that are candidates for + that mention. All items are packed into a single super-batch tensor so + the scorer is called exactly once. + + ``postprocess_sparse(packed_scores, meta)`` + Scatters the packed scores back into a full (B, seq_len, num_classes, 3) + output tensor. Unscored positions are filled with SPARSE_FILL. + + When ``_word_spans`` is None (dense / fallback path), ``forward`` is + byte-for-byte identical to the parent implementation — no overhead. + """ + + def __init__(self, config: Any, *args, **kwargs) -> None: + super().__init__(config, *args, **kwargs) + self._word_spans: List[List[Tuple[int, int]]] | None = None + self._label_indices: List[List[List[int]]] | None = None + + # ------------------------------------------------------------------ + # Sparse config management + # ------------------------------------------------------------------ + + def set_sparse_config( + self, + word_spans: List[List[Tuple[int, int]]], + label_indices: List[List[List[int]]], + ) -> None: + """Activate sparse mode for the next forward call. + + Args: + word_spans: ``word_spans[b][m] = (word_start, word_end)`` inclusive. + label_indices: ``label_indices[b][m]`` — indices into label_rep for mention m. + """ + self._word_spans = word_spans + self._label_indices = label_indices + + def clear_sparse_config(self) -> None: + """Deactivate sparse mode (restores dense behaviour).""" + self._word_spans = None + self._label_indices = None + + # ------------------------------------------------------------------ + # Sparse pre/post processing + # ------------------------------------------------------------------ + + def preprocess_sparse( + self, + token_rep: torch.Tensor, # (B, seq_len, H) + label_rep: torch.Tensor, # (B, num_classes, H) + ) -> Tuple[torch.Tensor | None, torch.Tensor | None, dict]: + """Pack sparse mentions into a single scorer call. + + Returns: + super_tok: ``(M_total, max_sl, H)`` — None if no valid spans. + super_lbl: ``(M_total, max_lc, H)`` — None if no valid spans. + meta: scatter metadata consumed by ``postprocess_sparse``. + """ + batch_size, seq_len, _H = token_rep.shape + device = token_rep.device + + all_tok: List[torch.Tensor] = [] + all_lbl: List[torch.Tensor] = [] + batch_meta: List[dict | None] = [] + + for b in range(batch_size): + ws_list = self._word_spans[b] if b < len(self._word_spans) else [] + li_list = self._label_indices[b] if b < len(self._label_indices) else [] + + if not ws_list or not li_list: + batch_meta.append(None) + continue + + M = len(ws_list) + ws_t = torch.tensor([s for s, _ in ws_list], device=device) # (M,) + we_t = torch.tensor([e for _, e in ws_list], device=device) # (M,) + slen_t = we_t - ws_t + 1 # (M,) + max_sl = int(slen_t.max()) + + n_lbl = torch.tensor([len(li) for li in li_list], device=device) # (M,) + max_lc = int(n_lbl.max()) + if max_lc == 0: + batch_meta.append(None) + continue + + # ── gather token reps: (M, max_sl, H) ──────────────────────── + t_range = torch.arange(max_sl, device=device) + tok_pos = (ws_t.unsqueeze(1) + t_range).clamp(0, seq_len - 1) # (M, max_sl) + tok = token_rep[b][tok_pos] # (M, max_sl, H) + tok_valid = t_range < slen_t.unsqueeze(1) # (M, max_sl) + tok = tok * tok_valid.unsqueeze(-1) # zero padding + + # ── gather label reps: (M, max_lc, H) ──────────────────────── + lbl_idx = torch.zeros(M, max_lc, device=device, dtype=torch.long) + for m, li in enumerate(li_list): + if li: + lbl_idx[m, : len(li)] = torch.tensor(li, device=device) + lbl = label_rep[b][lbl_idx] # (M, max_lc, H) + + all_tok.append(tok) + all_lbl.append(lbl) + batch_meta.append( + { + "b": b, + "M": M, + "tok_pos": tok_pos, # (M, max_sl) + "lbl_idx": lbl_idx, # (M, max_lc) + "tok_valid": tok_valid, # (M, max_sl) + "n_lbl": n_lbl, # (M,) + } + ) + + meta = { + "batch_meta": batch_meta, + "batch_size": batch_size, + "seq_len": seq_len, + "num_classes": label_rep.shape[1], + "device": device, + "dtype": token_rep.dtype, + } + + if not all_tok: + return None, None, meta + + # pad all items to the global (max_sl, max_lc) + global_sl = max(t.shape[1] for t in all_tok) + global_lc = max(l.shape[1] for l in all_lbl) + + super_tok = torch.cat([_pad_dim(t, global_sl, 1) for t in all_tok], dim=0) + super_lbl = torch.cat([_pad_dim(l, global_lc, 1) for l in all_lbl], dim=0) + return super_tok, super_lbl, meta + + def postprocess_sparse( + self, + packed_scores: torch.Tensor | None, # (M_total, global_sl, global_lc, 3) + meta: dict, + ) -> torch.Tensor: + """Scatter packed scores back into a full (B, seq_len, num_classes, 3) tensor. + + Unscored positions are pre-filled with SPARSE_FILL (≈ 0 after sigmoid). + """ + output = torch.full( + (meta["batch_size"], meta["seq_len"], meta["num_classes"], 3), + fill_value=SPARSE_FILL, + device=meta["device"], + dtype=meta["dtype"], + ) + if packed_scores is None: + return output + + offset = 0 + for bm in meta["batch_meta"]: + if bm is None: + continue + b, M = bm["b"], bm["M"] + tok_pos = bm["tok_pos"] # (M, local_sl) + lbl_idx = bm["lbl_idx"] # (M, local_lc) + tok_valid = bm["tok_valid"] # (M, local_sl) + n_lbl = bm["n_lbl"] # (M,) + + # slice back to local (potentially smaller) shapes + scores_b = packed_scores[ + offset : offset + M, + : tok_pos.shape[1], + : lbl_idx.shape[1], + ] # (M, local_sl, local_lc, 3) + + l_range = torch.arange(lbl_idx.shape[1], device=meta["device"]) + lbl_valid = l_range < n_lbl.unsqueeze(1) # (M, local_lc) + valid = tok_valid.unsqueeze(2) & lbl_valid.unsqueeze(1) # (M, local_sl, local_lc) + vm, vt, vl = valid.nonzero(as_tuple=True) + output[b, tok_pos[vm, vt], lbl_idx[vm, vl]] = scores_b[vm, vt, vl] + + offset += M + + return output + + # ------------------------------------------------------------------ + # Override forward + # ------------------------------------------------------------------ + + def forward( + self, + input_ids=None, + attention_mask=None, + labels_embeds=None, + labels_input_ids=None, + labels_attention_mask=None, + words_embedding=None, + mask=None, + span_idx=None, + span_mask=None, + span_labels=None, + prompts_embedding=None, + prompts_embedding_mask=None, + words_mask=None, + text_lengths=None, + labels=None, + threshold: float = 0.5, + **kwargs: Any, + ) -> GLiNERBaseOutput: + encoder_kwargs = { + key: kwargs[key] + for key in ("packing_config", "pair_attention_mask", "token_lengths", "word_lengths") + if key in kwargs + } + + prompts_embedding, prompts_embedding_mask, words_embedding, mask = self.get_representations( + input_ids, + attention_mask, + labels_embeds, + labels_input_ids, + labels_attention_mask, + text_lengths, + words_mask, + **encoder_kwargs, + ) + + if labels is not None: + words_embedding, mask = self._fit_length(words_embedding, mask, labels.shape[1]) + target_C = max(prompts_embedding.size(1), labels.size(-2)) + prompts_embedding, prompts_embedding_mask = self._fit_length( + prompts_embedding, prompts_embedding_mask, target_C + ) + + # ── scoring: sparse or dense ─────────────────────────────────────── + if self._word_spans is not None: + super_tok, super_lbl, meta = self.preprocess_sparse(words_embedding, prompts_embedding) + packed = self.scorer(super_tok, super_lbl) if super_tok is not None else None + scores = self.postprocess_sparse(packed, meta) + else: + scores = self.scorer(words_embedding, prompts_embedding) + + # ── optional span representations ───────────────────────────────── + if getattr(self.config, "represent_spans", False): + span_rep, span_idx, span_mask = self.get_span_representations( + scores, span_idx, span_mask, words_embedding, labels, threshold + ) + span_logits = torch.einsum("BND,BCD->BNC", span_rep, prompts_embedding) + else: + span_logits, span_idx, span_mask = None, None, None + + loss = None + if labels is not None: + loss = self.loss(scores, labels, prompts_embedding_mask, mask, **kwargs) + if span_labels is not None: + span_loss = self.loss( + span_logits, span_labels, prompts_embedding_mask, span_mask, **kwargs + ) + loss = self.config.token_loss_coef * loss + self.config.span_loss_coef * span_loss + + return GLiNERBaseOutput( + logits=scores, + loss=loss, + prompts_embedding=prompts_embedding, + prompts_embedding_mask=prompts_embedding_mask, + words_embedding=words_embedding, + mask=mask, + span_idx=span_idx, + span_logits=span_logits, + span_mask=span_mask, + ) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def build_word_spans( + gliner: Any, + text: str, + char_spans: List[Dict[str, int]], + span_label_indices: List[List[int]], +) -> Tuple[List[Tuple[int, int]], List[List[int]]]: + """Convert char-level spans to word-level spans. + + Uses GLiNER's own words_splitter so boundaries match exactly. + Spans that don't align to word boundaries are silently dropped. + + Args: + gliner: GLiNER instance (provides ``prepare_inputs``). + text: Input text. + char_spans: ``[{"start": int, "end": int}, ...]`` from L1. + span_label_indices: Per-span candidate indices aligned with char_spans. + + Returns: + ``(word_spans, label_indices)`` — only successfully matched spans. + """ + _, start_maps, end_maps = gliner.prepare_inputs([text]) + s2w = {char: word for word, char in enumerate(start_maps[0])} + e2w = {char: word for word, char in enumerate(end_maps[0])} + + word_spans: List[Tuple[int, int]] = [] + matched_lbl: List[List[int]] = [] + + for span, lbl_idxs in zip(char_spans, span_label_indices): + ws = s2w.get(span["start"]) + we = e2w.get(span["end"]) + if ws is not None and we is not None and we >= ws: + word_spans.append((ws, we)) + matched_lbl.append(list(lbl_idxs)) + + return word_spans, matched_lbl + + +# --------------------------------------------------------------------------- +# GLinkerModel +# --------------------------------------------------------------------------- + + +class GLinkerModel(BiEncoderTokenGLiNER): + """BiEncoderTokenGLiNER with sparse per-span scoring. + + Replaces the inner model with SparseTokenModel so that scoring is done only + over per-mention candidate labels when ``span_label_indices`` is provided. + + Usage:: + + model = GLinkerModel.from_pretrained("knowledgator/gliner-linker-large-v1.0") + + entities = model.predict_entities( + text, + labels, # flat union of all candidates + span_label_indices=[[0,1],[2,3]], # per-mention candidate indices + input_spans=[{"start":0,"end":5}, {"start":10,"end":18}], + ) + """ + + # ← This is the only line that replaces SparseScorer installation: + model_class = SparseTokenModel + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + def predict_entities( + self, + text: str, + labels: List[str], + span_label_indices: List[List[int]] | None = None, + input_spans: List[Dict[str, int]] | None = None, + flat_ner: bool = True, + threshold: float = 0.5, + span_extraction_threshold: float | None = None, + multi_label: bool = False, + return_class_probs: bool = False, + **kwargs: Any, + ) -> List[Dict[str, Any]]: + if span_label_indices is None: + return super().predict_entities( + text, + labels, + flat_ner=flat_ner, + threshold=threshold, + multi_label=multi_label, + return_class_probs=return_class_probs, + **kwargs, + ) + + word_spans, word_lbl_indices = build_word_spans(self, text, input_spans, span_label_indices) + return self.inference( + [text], + labels, + flat_ner=flat_ner, + threshold=threshold, + multi_label=multi_label, + return_class_probs=return_class_probs, + input_spans=[input_spans], + _word_spans=[word_spans], + _label_indices=[word_lbl_indices], + _span_extraction_threshold=span_extraction_threshold, + **kwargs, + )[0] + + def predict_with_embeds( + self, + text: str, + labels_embeddings: torch.Tensor, + labels: List[str], + span_label_indices: List[List[int]] | None = None, + input_spans: List[Dict[str, int]] | None = None, + flat_ner: bool = True, + threshold: float = 0.5, + span_extraction_threshold: float | None = None, + multi_label: bool = False, + return_class_probs: bool = False, + **kwargs: Any, + ) -> List[Dict[str, Any]]: + lbl_emb = labels_embeddings.to(self.device) + + if span_label_indices is None: + return self.inference( + [text], + labels, + flat_ner=flat_ner, + threshold=threshold, + multi_label=multi_label, + return_class_probs=return_class_probs, + labels_embeds=lbl_emb, + _span_extraction_threshold=span_extraction_threshold, + **kwargs, + )[0] + + word_spans, word_lbl_indices = build_word_spans(self, text, input_spans, span_label_indices) + return self.inference( + [text], + labels, + flat_ner=flat_ner, + threshold=threshold, + multi_label=multi_label, + return_class_probs=return_class_probs, + input_spans=[input_spans], + labels_embeds=lbl_emb, + _word_spans=[word_spans], + _label_indices=[word_lbl_indices], + _span_extraction_threshold=span_extraction_threshold, + **kwargs, + )[0] + + # ------------------------------------------------------------------ + # Batch processing — injects sparse config per batch + # ------------------------------------------------------------------ + + def _process_batches( + self, + data_loader, + threshold, + flat_ner, + multi_label, + packing_config=None, + return_class_probs=False, + word_input_spans=None, + _word_spans: List | None = None, + _label_indices: List | None = None, + _span_extraction_threshold: float | None = None, + **external_inputs, + ): + """Override _process_batches to handle sparse scoring and external inputs. + + Delegates to the parent ONLY when there is nothing special to do (no sparse + config, no external inputs like labels_embeds). Otherwise uses its own loop + because the installed GLiNER has a bug: it ignores ``external_inputs`` when + ``packing_config is None`` (merges only when packing_config is set). + + ``_span_extraction_threshold`` is passed to ``model.forward`` (controls BIO + span extraction), while ``threshold`` is used by ``decoder.decode`` (filters + final span_logits). When None, ``threshold`` is used for both. + """ + if _word_spans is None and not external_inputs and _span_extraction_threshold is None: + return super()._process_batches( + data_loader, + threshold, + flat_ner, + multi_label, + packing_config=packing_config, + return_class_probs=return_class_probs, + word_input_spans=word_input_spans, + ) + + # BIO extraction uses a separate (typically lower) threshold if provided. + bio_threshold = ( + _span_extraction_threshold if _span_extraction_threshold is not None else threshold + ) + + # Custom loop: handles sparse config AND external inputs (e.g. labels_embeds). + outputs = [] + device = self.device + batch_offset = 0 + + for batch in data_loader: + if not getattr(self, "onnx_model", False): + batch = { + k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in batch.items() + } + + current_batch_size = len(batch["tokens"]) + model_inputs = {**batch, **external_inputs} + if packing_config is not None: + model_inputs["packing_config"] = packing_config + + if _word_spans is not None: + batch_ws = _word_spans[batch_offset : batch_offset + current_batch_size] + batch_li = _label_indices[batch_offset : batch_offset + current_batch_size] + self.model.set_sparse_config(batch_ws, batch_li) + + try: + with torch.inference_mode(): + model_output = self.model(**model_inputs, threshold=bio_threshold) + finally: + if _word_spans is not None: + self.model.clear_sparse_config() + + model_logits = model_output[0] + if not isinstance(model_logits, torch.Tensor): + model_logits = torch.from_numpy(model_logits) + + batch_input_spans = None + if word_input_spans is not None: + batch_input_spans = word_input_spans[ + batch_offset : batch_offset + current_batch_size + ] + + decoded = self.decoder.decode( + batch["tokens"], + batch["id_to_classes"], + model_logits, + span_idx=model_output.span_idx, + span_mask=model_output.span_mask, + span_logits=model_output.span_logits, + flat_ner=flat_ner, + threshold=threshold, + multi_label=multi_label, + return_class_probs=return_class_probs, + input_spans=batch_input_spans, + ) + outputs.extend(decoded) + batch_offset += current_batch_size + + return outputs diff --git a/src/glinker/l2/component.py b/src/glinker/l2/component.py index e1cf05e..c1b72f4 100644 --- a/src/glinker/l2/component.py +++ b/src/glinker/l2/component.py @@ -239,7 +239,21 @@ def is_available(self) -> bool: class RedisLayer(DatabaseLayer): - """Redis cache layer.""" + """Redis cache layer with optimized embedding storage. + + Storage structure: + - entity:{label} -> entity data without embedding (fast lookup) + - entity:{alias} -> entity data without embedding (fast lookup) + - entity:emb:{entity_id} -> {embedding, embedding_model_id} (no duplication) + + Benefits: + - Embeddings stored once per entity (not duplicated across aliases) + - update_embeddings() is O(M) instead of O(NxM) where: + - M = entities to update + - N = total keys in Redis (entities x aliases) + - Reduced memory usage for large entity databases + - Faster batch embedding updates + """ def _setup(self): self.client = redis.Redis( @@ -288,25 +302,58 @@ def search_fuzzy(self, query: str) -> List[DatabaseRecord]: return [] def write_cache(self, key: str, records: List[DatabaseRecord], ttl: int): + """Write search results to cache. + + Optimized: embeddings stored separately to avoid duplication. + """ key = self.normalize_query(key) cache_key = f"entity:{key}" try: - data = json.dumps([r.dict() for r in records]) - self.client.setex(cache_key, ttl, data) + # Store entity data WITHOUT embeddings + records_data = [] + pipe = self.client.pipeline() + + for r in records: + r_dict = r.dict() + embedding = r_dict.pop("embedding", None) + embedding_model_id = r_dict.pop("embedding_model_id", None) + records_data.append(r_dict) + + # Store embedding separately if present + if embedding is not None: + emb_key = f"entity:emb:{r.entity_id}" + emb_data = json.dumps( + {"embedding": embedding, "embedding_model_id": embedding_model_id} + ) + pipe.setex(emb_key, ttl, emb_data) + + # Store main cache data + data = json.dumps(records_data) + pipe.setex(cache_key, ttl, data) + pipe.execute() + except Exception as e: print(f"[ERROR Redis] Write error: {e}") def load_bulk( self, entities: List[DatabaseRecord], overwrite: bool = False, batch_size: int = 1000 ) -> int: - """Bulk load to Redis.""" + """Bulk load to Redis. + + Optimized: embeddings stored separately to avoid duplication across aliases. + Structure: + - entity:{label/alias} -> entity data WITHOUT embedding + - entity:emb:{entity_id} -> {embedding, model_id} + """ count = 0 pipe = self.client.pipeline() for entity in entities: - # Prepare data + # Prepare data WITHOUT embedding (store separately) entity_data = entity.dict() + embedding = entity_data.pop("embedding", None) + embedding_model_id = entity_data.pop("embedding_model_id", None) data_json = json.dumps(entity_data) # Store by label @@ -321,6 +368,14 @@ def load_bulk( if overwrite or not self.client.exists(alias_key): pipe.setex(alias_key, self.ttl, data_json) + # Store embedding SEPARATELY (no duplication) + if embedding is not None: + emb_key = f"entity:emb:{entity.entity_id}" + emb_data = json.dumps( + {"embedding": embedding, "embedding_model_id": embedding_model_id} + ) + pipe.setex(emb_key, self.ttl, emb_data) + # Execute in batches if len(pipe) >= batch_size: pipe.execute() @@ -333,20 +388,36 @@ def load_bulk( return count def clear(self): - """Clear all entity keys.""" + """Clear all entity keys (including embeddings).""" for key in self.client.scan_iter(match="entity:*"): self.client.delete(key) + # Note: entity:emb:* keys are also matched by entity:* pattern def count(self) -> int: """Count entity keys.""" return sum(1 for _ in self.client.scan_iter(match="entity:*")) - def get_all_entities(self) -> List[DatabaseRecord]: - """Get all entities from Redis (scans all entity:* keys).""" + def get_all_entities(self, include_embeddings: bool = False) -> List[DatabaseRecord]: + """Get all entities from Redis (scans all entity:* keys). + + Args: + include_embeddings: If True, also fetch embeddings (slower) + + Note: Skips entity:emb:* keys as they're fetched separately if needed. + """ entities = [] seen_ids = set() for key in self.client.scan_iter(match="entity:*"): + # Skip embedding keys - they're handled separately + if isinstance(key, bytes): + key_str = key.decode("utf-8") + else: + key_str = key + + if key_str.startswith("entity:emb:"): + continue + try: data = self.client.get(key) if data: @@ -368,48 +439,85 @@ def get_all_entities(self) -> List[DatabaseRecord]: except Exception: continue + # Optionally fetch embeddings in batch + if include_embeddings and entities: + entity_ids = [e.entity_id for e in entities] + embeddings_map = self.get_embeddings_batch(entity_ids) + + for entity in entities: + if entity.entity_id in embeddings_map: + emb_data = embeddings_map[entity.entity_id] + entity.embedding = emb_data.get("embedding") + entity.embedding_model_id = emb_data.get("embedding_model_id") + return entities def update_embeddings( self, entity_ids: List[str], embeddings: List[List[float]], model_id: str ) -> int: - """Update embeddings in Redis entities.""" - count = 0 - id_to_embedding = dict(zip(entity_ids, embeddings)) + """Update embeddings in Redis entities. - for key in self.client.scan_iter(match="entity:*"): - try: - data = self.client.get(key) - if not data: - continue + Optimized: O(M) direct access instead of O(NxM) full scan. + Embeddings stored separately at entity:emb:{entity_id}. + """ + if not entity_ids or not embeddings: + return 0 - if isinstance(data, bytes): - data = data.decode("utf-8") + pipe = self.client.pipeline() + count = 0 - record_data = json.loads(data) - updated = False - - if isinstance(record_data, dict): - if record_data.get("entity_id") in id_to_embedding: - record_data["embedding"] = id_to_embedding[record_data["entity_id"]] - record_data["embedding_model_id"] = model_id - updated = True - elif isinstance(record_data, list): - for r in record_data: - if r.get("entity_id") in id_to_embedding: - r["embedding"] = id_to_embedding[r["entity_id"]] - r["embedding_model_id"] = model_id - updated = True - - if updated: - self.client.setex(key, self.ttl, json.dumps(record_data)) - count += 1 + for entity_id, embedding in zip(entity_ids, embeddings): + emb_key = f"entity:emb:{entity_id}" + emb_data = json.dumps({"embedding": embedding, "embedding_model_id": model_id}) + pipe.setex(emb_key, self.ttl, emb_data) + count += 1 - except Exception: - continue + # Execute in batches of 1000 + if count % 1000 == 0: + pipe.execute() + pipe = self.client.pipeline() + + # Execute remaining + if len(pipe) > 0: + pipe.execute() return count + def get_embeddings_batch(self, entity_ids: List[str]) -> Dict[str, Dict[str, Any]]: + """Get embeddings for multiple entities in batch. + + Args: + entity_ids: List of entity IDs + + Returns: + Dict mapping entity_id -> {embedding: List[float], embedding_model_id: str} + Missing entities are omitted from result. + """ + if not entity_ids: + return {} + + # Batch GET using pipeline + pipe = self.client.pipeline() + for entity_id in entity_ids: + emb_key = f"entity:emb:{entity_id}" + pipe.get(emb_key) + + results = pipe.execute() + + # Parse results + embeddings_map = {} + for entity_id, data in zip(entity_ids, results): + if data: + try: + if isinstance(data, bytes): + data = data.decode("utf-8") + emb_data = json.loads(data) + embeddings_map[entity_id] = emb_data + except Exception: + continue + + return embeddings_map + def is_available(self) -> bool: try: self.client.ping() @@ -427,6 +535,30 @@ def _setup(self): ) self.index_name = self.config.config["index_name"] self.popularity_boost = self.config.config.get("popularity_boost", False) + self._ensure_index() + + def _ensure_index(self): + """Create index with explicit mapping so dense_vector is stored in _source.""" + if self.client.indices.exists(index=self.index_name): + return + self.client.indices.create( + index=self.index_name, + body={ + "mappings": { + "properties": { + "label": {"type": "text", "analyzer": "english"}, + "aliases": {"type": "text", "analyzer": "english"}, + "description": {"type": "text", "analyzer": "english"}, + "entity_id": {"type": "keyword"}, + "entity_type": {"type": "keyword"}, + "popularity": {"type": "long"}, + "embedding_model_id": {"type": "keyword"}, + # index:false keeps the vector in _source without building an HNSW index + "embedding": {"type": "dense_vector", "index": False}, + } + } + }, + ) def _build_query(self, match_query: dict) -> dict: """Wrap a match query with optional popularity boosting. @@ -438,8 +570,11 @@ def _build_query(self, match_query: dict) -> dict: Uses ln2p (not ln1p) to avoid zeroing out entities with popularity=0, since ln(2+0)=0.69 while ln(1+0)=0. """ + # dense_vector is excluded from _source by default in ES 8+; request it explicitly + _source = {"includes": ["*", "embedding", "embedding_model_id"]} + if not self.popularity_boost: - return {"query": match_query, "size": 50} + return {"query": match_query, "size": 50, "_source": _source} return { "query": { @@ -450,6 +585,7 @@ def _build_query(self, match_query: dict) -> dict: } }, "size": 50, + "_source": _source, } def search(self, query: str) -> List[DatabaseRecord]: @@ -491,6 +627,73 @@ def search_fuzzy(self, query: str) -> List[DatabaseRecord]: print(f"[ERROR ES] Fuzzy error: {e}") return [] + def batch_search(self, queries: List[str], fuzzy: bool = False) -> List[List[DatabaseRecord]]: + """ + Batch search using Elasticsearch msearch API. + + Args: + queries: List of query strings + fuzzy: Whether to use fuzzy matching + + Returns: + List of result lists, one per query (same order as input) + """ + if not queries: + return [] + + # Normalize all queries + normalized_queries = [self.normalize_query(q) for q in queries] + + # Build msearch request body + request_body = [] + for query in normalized_queries: + # Header for each search + request_body.append({"index": self.index_name}) + + # Body for each search + if fuzzy and self.supports_fuzzy(): + fuzzy_distance = self.fuzzy_config.max_distance + match_query = { + "multi_match": { + "query": query, + "fields": ["label^2", "aliases^1.5", "description"], + "fuzziness": fuzzy_distance, + "prefix_length": self.fuzzy_config.prefix_length, + "max_expansions": 50, + } + } + else: + match_query = { + "multi_match": { + "query": query, + "fields": ["label^2", "aliases^1.5", "description"], + "type": "best_fields", + } + } + + body = self._build_query(match_query) + request_body.append(body) + + try: + # Execute msearch + response = self.client.msearch(body=request_body) + + # Process responses + results = [] + for resp in response["responses"]: + if "error" in resp: + print(f"[ERROR ES] Batch search error: {resp['error']}") + results.append([]) + else: + hits = resp.get("hits", {}).get("hits", []) + results.append(self._process_hits(hits)) + + return results + + except Exception as e: + print(f"[ERROR ES] Batch search error: {e}") + return [[] for _ in queries] + def _process_hits(self, hits: List[Dict]) -> List[DatabaseRecord]: records = [] for hit in hits: @@ -961,6 +1164,78 @@ def search(self, mention: str) -> List[DatabaseRecord]: return results + def batch_search(self, mentions: List[str]) -> List[List[DatabaseRecord]]: + """ + Batch search through layers with fallback. + + Args: + mentions: List of mention strings to search + + Returns: + List of result lists, one per mention (same order as input) + """ + if not mentions: + return [] + + # Try each layer in priority order + for layer in self.layers: + if not layer.is_available(): + continue + + # Check if layer supports batch_search + if not hasattr(layer, "batch_search"): + # Fallback to individual searches + results = [] + for mention in mentions: + layer_results = [] + for mode in layer.config.search_mode: + if mode == "exact": + layer_results.extend(layer.search(mention)) + elif mode == "fuzzy": + if layer.supports_fuzzy(): + layer_results.extend(layer.search_fuzzy(mention)) + + if layer_results: + layer_results = self.deduplicate_candidates(layer_results) + results.append(layer_results) + + # Check if we got any results + if any(results): + return results + else: + # Use batch search + all_results = [] + + for mode in layer.config.search_mode: + if mode == "exact": + batch_results = layer.batch_search(mentions, fuzzy=False) + all_results.append(batch_results) + elif mode == "fuzzy": + if layer.supports_fuzzy(): + batch_results = layer.batch_search(mentions, fuzzy=True) + all_results.append(batch_results) + + # Merge results from different modes + if all_results: + merged = [] + for i in range(len(mentions)): + mention_results = [] + for mode_results in all_results: + if i < len(mode_results): + mention_results.extend(mode_results[i]) + + if mention_results: + mention_results = self.deduplicate_candidates(mention_results) + merged.append(mention_results) + + # Check if we got any results + if any(merged): + # TODO: Implement batch cache write + return merged + + # No results from any layer + return [[] for _ in mentions] + def _cache_write(self, query: str, results: List[DatabaseRecord], source_layer: DatabaseLayer): """Write results to upper layers (higher priority = checked earlier).""" for layer in self.layers: diff --git a/src/glinker/l2/models.py b/src/glinker/l2/models.py index 44be503..66bcaf2 100644 --- a/src/glinker/l2/models.py +++ b/src/glinker/l2/models.py @@ -100,4 +100,13 @@ class L2Input(BaseInput): class L2Output(BaseOutput): """L2 processor output.""" - candidates: List[List[DatabaseRecord]] = Field(..., description="Candidates per mention/group") + candidates: List[List[DatabaseRecord]] = Field( + ..., description="Flat candidates per text — candidates[text_idx][candidate] (used by L4)" + ) + per_span_candidates: List[List[List[DatabaseRecord]]] | None = Field( + None, + description=( + "Per-span candidates: per_span_candidates[text_idx][span_idx][candidate_idx]. " + "Populated when mentions are nested (List[List[entity]] from L1). Used by L3." + ), + ) diff --git a/src/glinker/l2/processor.py b/src/glinker/l2/processor.py index 3b8f4da..8569792 100644 --- a/src/glinker/l2/processor.py +++ b/src/glinker/l2/processor.py @@ -90,31 +90,99 @@ def __call__( # Check if mentions is nested (list of lists - one per text) if mentions and isinstance(mentions[0], (list, tuple)): # Nested structure: [[entities_text1], [entities_text2], ...] - all_candidates = [] - - for text_entities in mentions: - text_candidates = [] - - for entity in text_entities: - # Extract text from L1Entity or dict - mention_text = self._extract_mention_text(entity) - - # Search candidates for this mention - candidates = self._execute_pipeline(mention_text, self.pipeline) - text_candidates.extend(candidates) - - all_candidates.append(text_candidates) - - return L2Output(candidates=all_candidates) + # We build both: + # all_flat[text_idx] = flat list of all candidates (for L4 / backward compat) + # all_per_span[text_idx][span_idx] = per-span lists (for L3 sparse scoring) + all_flat = [] + all_per_span = [] + + # OPTIMIZATION: Use batch search if available + if hasattr(self.component, "batch_search"): + # Flatten mentions and track structure + flat_mentions = [] + mention_counts = [] + + for text_entities in mentions: + mention_texts = [self._extract_mention_text(e) for e in text_entities] + flat_mentions.extend(mention_texts) + mention_counts.append(len(mention_texts)) + + # Batch search all mentions at once + if flat_mentions: + batch_results = self.component.batch_search(flat_mentions) + + # Apply post-processing pipeline to each result + processed_results = [] + for result in batch_results: + processed = result + for step_name, step_kwargs in self.pipeline: + if step_name != "search": + method = getattr(self.component, step_name) + processed = method(processed, **step_kwargs) + processed_results.append(processed) + + # Reconstruct: flat (extend) for L4, per-span (append) for L3 + idx = 0 + for count in mention_counts: + text_flat: list = [] + text_per_span: list = [] + for _ in range(count): + if idx < len(processed_results): + span_cands = processed_results[idx] + text_flat.extend(span_cands) + text_per_span.append(span_cands) + idx += 1 + all_flat.append(text_flat) + all_per_span.append(text_per_span) + else: + all_flat = [[] for _ in mentions] + all_per_span = [[] for _ in mentions] + else: + # Fallback: individual searches + for text_entities in mentions: + text_flat = [] + text_per_span = [] + for entity in text_entities: + mention_text = self._extract_mention_text(entity) + span_cands = self._execute_pipeline(mention_text, self.pipeline) + text_flat.extend(span_cands) + text_per_span.append(span_cands) + all_flat.append(text_flat) + all_per_span.append(text_per_span) + + return L2Output(candidates=all_flat, per_span_candidates=all_per_span) # Flat structure: ["mention1", "mention2", ...] else: - all_candidates = [] + # OPTIMIZATION: Use batch search if available + if hasattr(self.component, "batch_search"): + # Extract mention texts + mention_texts = [self._extract_mention_text(m) for m in mentions] + + # Batch search all at once + if mention_texts: + batch_results = self.component.batch_search(mention_texts) + + # Apply post-processing pipeline to each result + all_candidates = [] + for result in batch_results: + # Apply pipeline steps (filter, deduplicate, sort, limit) + processed = result + for step_name, step_kwargs in self.pipeline: + if step_name != "search": # Skip search, already done + method = getattr(self.component, step_name) + processed = method(processed, **step_kwargs) + all_candidates.append(processed) + else: + all_candidates = [] + else: + # Fallback: individual searches + all_candidates = [] - for mention in mentions: - mention_text = self._extract_mention_text(mention) - candidates = self._execute_pipeline(mention_text, self.pipeline) - all_candidates.append(candidates) + for mention in mentions: + mention_text = self._extract_mention_text(mention) + candidates = self._execute_pipeline(mention_text, self.pipeline) + all_candidates.append(candidates) if structure: grouped = self._group_by_structure(all_candidates, structure) diff --git a/src/glinker/l3/component.py b/src/glinker/l3/component.py index 00ed15a..a85bd58 100644 --- a/src/glinker/l3/component.py +++ b/src/glinker/l3/component.py @@ -1,9 +1,9 @@ from typing import List import torch -from gliner import GLiNER from glinker.core.base import BaseComponent +from glinker.core.gliner_model import GLinkerModel from .models import L3Config, L3Entity @@ -12,22 +12,24 @@ class L3Component(BaseComponent[L3Config]): """GLiNER-based entity linking component.""" 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 = GLinkerModel.from_pretrained( + self.config.model_name, + token=self.config.token, + max_length=self.config.max_length, ) self.model.to(self.config.device) - # Fix labels tokenizer max_length for BiEncoder models - # Some models have model_max_length not properly set (> 10^18) + # Fix labels tokenizer max_length for BiEncoder models. + # label_max_length takes priority; falls back to max_length. + _lbl_len = self.config.label_max_length or self.config.max_length if ( - self.config.max_length is not None + _lbl_len is not None and hasattr(self.model, "data_processor") and hasattr(self.model.data_processor, "labels_tokenizer") ): tok = self.model.data_processor.labels_tokenizer - if tok.model_max_length > 100000: - tok.model_max_length = self.config.max_length + if tok.model_max_length > 100_000: + tok.model_max_length = _lbl_len @property def device(self): @@ -35,7 +37,6 @@ def device(self): @property def supports_precomputed_embeddings(self) -> bool: - """Check if model supports precomputed embeddings (BiEncoder).""" return hasattr(self.model, "encode_labels") and self.model.config.labels_encoder is not None def get_available_methods(self) -> List[str]: @@ -49,125 +50,85 @@ def get_available_methods(self) -> List[str]: ] def encode_labels(self, labels: List[str], batch_size: int = 32) -> torch.Tensor: - """ - Encode labels using GLiNER's native label encoder. - - Args: - labels: List of label strings to encode - batch_size: Batch size for encoding - - Returns: - Tensor of shape (num_labels, hidden_size) - - Raises: - NotImplementedError: If model doesn't support label encoding (not BiEncoder) - """ if not self.supports_precomputed_embeddings: raise NotImplementedError( f"Model {self.config.model_name} doesn't support label precomputation. " "Only BiEncoder models support this feature." ) - return self.model.encode_labels(labels, batch_size=batch_size) - def predict_with_embeddings( + @property + def _span_extraction_threshold(self) -> float: + """BIO extraction threshold — falls back to threshold if not set.""" + return self.config.span_extraction_threshold or self.config.threshold + + def predict_entities( self, text: str, labels: List[str], - embeddings: torch.Tensor, - input_spans: List[List[dict]] | None = None, + input_spans: List[dict] | None = None, + span_label_indices: List[List[int]] | None = None, ) -> List[L3Entity]: - """ - Predict entities using pre-computed label embeddings. - - Args: - text: Input text - labels: List of label strings (for output mapping) - embeddings: Pre-computed embeddings tensor (num_labels, hidden_size) - input_spans: Optional list of span dicts with 'start' and 'end' keys - to constrain prediction to specific spans from L1 + """Predict entities with optional sparse scoring. - Returns: - List of L3Entity predictions + When both *input_spans* and *span_label_indices* are provided each + mention is scored only against its own candidates (sparse path). + Otherwise falls back to standard dense GLiNER scoring. """ - if not self.supports_precomputed_embeddings: - # Fallback to regular prediction - return self.predict_entities(text, labels, input_spans=input_spans) - - kwargs = { - "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_with_embeds(text, embeddings, 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 - ] + if not labels: + return [] + entities = self.model.predict_entities( + text, + labels, + span_label_indices=span_label_indices, + input_spans=input_spans, + threshold=self.config.threshold, + span_extraction_threshold=self._span_extraction_threshold, + flat_ner=self.config.flat_ner, + multi_label=self.config.multi_label, + return_class_probs=True, + ) + return [self._to_entity(e) for e in entities] - def predict_entities( - self, text: str, labels: List[str], input_spans: List[List[dict]] | None = None + def predict_with_embeddings( + self, + text: str, + labels: List[str], + embeddings: torch.Tensor, + input_spans: List[dict] | None = None, + span_label_indices: List[List[int]] | None = None, ) -> List[L3Entity]: - """Predict entities using GLiNER. - - Args: - text: Input text - labels: List of label strings - input_spans: Optional list of span dicts with 'start' and 'end' keys - to constrain prediction to specific spans from L1 - """ + """Predict from pre-computed embeddings with optional sparse scoring.""" if not labels: return [] - - kwargs = { - "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 - ] + if not self.supports_precomputed_embeddings: + return self.predict_entities(text, labels, input_spans, span_label_indices) + + entities = self.model.predict_with_embeds( + text, + embeddings, + labels, + span_label_indices=span_label_indices, + input_spans=input_spans, + threshold=self.config.threshold, + span_extraction_threshold=self._span_extraction_threshold, + flat_ner=self.config.flat_ner, + multi_label=self.config.multi_label, + return_class_probs=True, + ) + return [self._to_entity(e) for e in entities] def filter_by_score( self, entities: List[L3Entity], threshold: float | None = 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.""" - seen = set() + seen: set = set() unique = [] for entity in entities: key = (entity.text, entity.start, entity.end) @@ -175,3 +136,14 @@ def deduplicate_entities(self, entities: List[L3Entity]) -> List[L3Entity]: unique.append(entity) seen.add(key) return unique + + @staticmethod + def _to_entity(e: dict) -> L3Entity: + return L3Entity( + text=e["text"], + label=e["label"], + start=e["start"], + end=e["end"], + score=e["score"], + class_probs=e.get("class_probs"), + ) diff --git a/src/glinker/l3/models.py b/src/glinker/l3/models.py index eb53a1f..3135f5c 100644 --- a/src/glinker/l3/models.py +++ b/src/glinker/l3/models.py @@ -19,12 +19,42 @@ class L3Config(BaseConfig): True, description="Use precomputed embeddings from L2 candidates if available" ) cache_embeddings: bool = Field(False, description="Cache computed embeddings back to L2") + span_extraction_threshold: float | None = Field( + None, + description=( + "Threshold for BIO span extraction inside model.forward. " + "Lower → more candidate spans constructed → higher recall. " + "When None, falls back to `threshold`. " + "Useful to set low (0.1-0.2) so final `threshold` filters span_logits " + "without also cutting span candidates." + ), + ) max_length: int = Field( None, - description="Maximum sequence length for tokenization. Passed to GLiNER.from_pretrained.", + description="Maximum sequence length for text tokenization. Passed to GLiNER.from_pretrained.", + ) + label_max_length: int = Field( + None, + description=( + "Maximum token length for label encoding. " + "Overrides max_length for the labels tokenizer only. " + "Useful when labels are short descriptions and you want to save memory." + ), ) +class L3LLMConfig(BaseConfig): + provider: str = Field("gemini", description="LLM provider: 'gemini' or 'openai'") + model_name: str = Field("gemini-2.0-flash") + api_key: str | None = Field(None, description="API key (falls back to env var)") + top_k: int = Field(5, description="Max candidates shown to the LLM per mention") + threshold: float = Field(0.0) + batch_size: int = Field(10, description="Mentions per LLM request") + max_retries: int = Field(3) + retry_delay: float = Field(1.0) + temperature: float = Field(0.0) + + # TODO replace candidates with labels class L3Input(BaseInput): texts: List[str] = Field(...) diff --git a/src/glinker/l3/processor.py b/src/glinker/l3/processor.py index 01e4b8d..44fa4a1 100644 --- a/src/glinker/l3/processor.py +++ b/src/glinker/l3/processor.py @@ -1,3 +1,10 @@ +from __future__ import annotations + +import os +import re +import json +import time +import logging from typing import Any, List import torch @@ -5,9 +12,11 @@ from glinker.core.base import BaseProcessor from glinker.core.registry import processor_registry -from .models import L3Input, L3Config, L3Entity, L3Output +from .models import L3Input, L3Config, L3Entity, L3Output, L3LLMConfig from .component import L3Component +logger = logging.getLogger(__name__) + class L3Processor(BaseProcessor[L3Config, L3Input, L3Output]): """GLiNER entity linking processor.""" @@ -24,146 +33,140 @@ def __init__( 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"] if isinstance(e, dict) else e.start, - "end": e["end"] if isinstance(e, dict) else e.end, - } - for e in l1_entities_for_text - ] - return [spans] + return [("predict_entities", {}), ("sort_by_position", {})] def __call__( self, - texts: List[str] | None = None, - candidates: List[List[Any]] | None = None, + texts: List[str], + candidates: List[List[List[Any]]], l1_entities: List[List[Any]] | None = None, input_data: L3Input = None, ) -> L3Output: - """Process texts with candidate labels. + """Process texts with per-span 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 + texts: List of input texts. + candidates: ``candidates[text_idx][span_idx]`` — L2 candidates per span. + l1_entities: ``l1_entities[text_idx]`` — L1 entities defining the spans. """ - # 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 = {} + for idx, (text, per_span_candidates) in enumerate(zip(texts, candidates)): + text_l1 = l1_entities[idx] if l1_entities and idx < len(l1_entities) else [] + + # input_spans from L1 — char-level, aligned with per_span_candidates + input_spans = [ + { + "start": e["start"] if isinstance(e, dict) else e.start, + "end": e["end"] if isinstance(e, dict) else e.end, + } + for e in text_l1 + ] + + # Build flat labels + per-span indices from L2 per-span candidates + labels, span_label_indices, label_to_candidate = self._build_sparse_inputs( + per_span_candidates + ) - shared_use_precomputed = ( + if not labels: + all_entities.append([]) + continue + + 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 + and self._can_use_precomputed( + [c for span in per_span_candidates for c in span], + 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) + if use_precomputed: + flat_candidates = [c for span in per_span_candidates for c in span] + embeddings = self._get_embeddings_tensor( + flat_candidates, labels, 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 + text, labels, embeddings, input_spans, span_label_indices ) + elif self.component.supports_precomputed_embeddings: + # Encode in chunks (encode_labels defaults to batch_size=32) so we + # never push thousands of labels through DeBERTa in a single pass → OOM. + embeddings = self.component.encode_labels(labels) + entities = self.component.predict_with_embeddings( + text, labels, embeddings, input_spans, span_label_indices + ) + if self.config.cache_embeddings: + flat_candidates = [c for span in per_span_candidates for c in span] + # Pass the tensor we just computed — skip a second encode_labels call. + self._cache_embeddings_tensor( + flat_candidates, labels, label_to_candidate, embeddings + ) 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) + entities = self.component.predict_entities( + text, labels, input_spans, span_label_indices + ) - # Apply rest of pipeline + # Apply rest of pipeline (sort, dedup, etc.) 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) + flat_candidates = [c for span in per_span_candidates for c in span] + entities = self._rank_entities(entities, flat_candidates) all_entities.append(entities) return L3Output(entities=all_entities) + def _build_sparse_inputs( + self, + per_span_candidates: List[List[Any]], + ) -> tuple[List[str], List[List[int]], dict]: + """Build flat labels list and per-span label indices from L2 per-span candidates. + + Args: + per_span_candidates: ``per_span_candidates[span_idx]`` = list of candidates. + + Returns: + ``(labels, span_label_indices, label_to_candidate)`` + - labels: flat deduplicated label strings + - span_label_indices: for each span, indices into labels + - label_to_candidate: label string → candidate object + """ + template = self.schema.get("template", "{label}") + label_to_idx: dict = {} + labels: List[str] = [] + label_to_candidate: dict = {} + span_label_indices: List[List[int]] = [] + + for span_cands in per_span_candidates: + idxs: List[int] = [] + for cand in span_cands: + lbl = self._format_label(cand, template) + if lbl not in label_to_idx: + label_to_idx[lbl] = len(labels) + labels.append(lbl) + label_to_candidate[lbl] = cand + idxs.append(label_to_idx[lbl]) + span_label_indices.append(idxs) + + return labels, span_label_indices, label_to_candidate + + def _format_label(self, candidate: Any, template: str) -> str: + """Format a candidate into a label string using the schema template.""" + try: + if hasattr(candidate, "model_dump"): + return template.format(**candidate.model_dump()) + if hasattr(candidate, "dict"): + return template.format(**candidate.dict()) + if isinstance(candidate, dict): + return template.format(**candidate) + except KeyError: + pass + return getattr(candidate, "label", str(candidate)) + def _can_use_precomputed(self, candidates: List[Any], label_to_candidate: dict) -> bool: """Check if all candidates have compatible precomputed embeddings.""" if not candidates: @@ -219,18 +222,7 @@ def _cache_embeddings(self, candidates: List[Any], labels: List[str], label_to_c # 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) + compute_labels = [self._format_label(c, template) for c in to_compute] # Encode labels embeddings = self.component.encode_labels(compute_labels) @@ -244,58 +236,46 @@ def _cache_embeddings(self, candidates: List[Any], labels: List[str], label_to_c ) 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 _cache_embeddings_tensor( + self, + candidates: List[Any], + labels: List[str], + label_to_candidate: dict, + embeddings: torch.Tensor, + ): + """Cache pre-computed embeddings for candidates that don't have them yet. - def _create_gliner_labels_with_mapping(self, candidates: List[Any]) -> tuple: + Unlike :meth:`_cache_embeddings`, accepts an already-encoded tensor so + we skip a redundant ``encode_labels`` call when the caller just encoded. + ``embeddings[i]`` must correspond to ``labels[i]``. """ - Create GLiNER labels using schema template and return label->candidate mapping. + if not self._l2_processor: + return - Returns: - tuple: (labels: List[str], label_to_candidate: dict) - """ + label_to_emb_idx = {lbl: i for i, lbl in enumerate(labels)} template = self.schema.get("template", "{label}") - labels = [] - label_to_candidate = {} - seen = set() + + to_update_ids: List[Any] = [] + to_update_vecs: List[Any] = [] 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 + if getattr(candidate, "embedding", None): + continue # already cached + lbl = self._format_label(candidate, template) + emb_idx = label_to_emb_idx.get(lbl) + if emb_idx is None: + continue + to_update_ids.append(candidate.entity_id) + to_update_vecs.append(embeddings[emb_idx].tolist()) + + if not to_update_ids: + return - 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 + if hasattr(self._l2_processor, "component"): + for layer in self._l2_processor.component.layers: + if layer.is_available(): + layer.update_embeddings(to_update_ids, to_update_vecs, self.config.model_name) + break def _rank_entities(self, entities: List[L3Entity], candidates: List[Any]) -> List[L3Entity]: """Re-rank entities using multiple scoring factors.""" @@ -341,3 +321,197 @@ def create_l3_processor(config_dict: dict, pipeline: list | None = None) -> L3Pr config = L3Config(**config_dict) component = L3Component(config) return L3Processor(config, component, pipeline) + + +# ── LLM-based L3 ───────────────────────────────────────────────────────────── + +_SYSTEM_PROMPT = """\ +You are an entity disambiguation assistant. +For each mention you will receive: + - The sentence with the mention wrapped in [[ ]]. + - A numbered list of candidate entities (label: description). + +Reply with ONLY a JSON array, one entry per mention, in the same order. +Each entry must be either: + - an integer index (0-based) of the correct candidate, or + - null if none of the candidates fits. + +Example reply: [0, null, 2] +Do NOT add any explanation or extra text.""" + + +def _mark_mention(text: str, start: int, end: int) -> str: + return text[:start] + "[[" + text[start:end] + "]]" + text[end:] + + +def _build_prompt(batch: list[dict]) -> str: + lines = [] + for i, item in enumerate(batch): + lines.append(f"Mention {i}: {item['context']}") + for j, c in enumerate(item["candidates"]): + desc = c.get("description") or "" + label = c.get("label", "") + lines.append(f" {j}. {label}: {desc}" if desc else f" {j}. {label}") + lines.append("") + return "\n".join(lines) + + +def _parse_llm_response(raw: str, expected: int) -> list[int | None]: + match = re.search(r"\[.*?\]", raw, re.DOTALL) + if not match: + logger.warning(f"No JSON array in LLM response: {raw!r}") + return [None] * expected + try: + parsed = json.loads(match.group()) + except json.JSONDecodeError as e: + logger.warning(f"JSON parse error: {e}. Raw: {raw!r}") + return [None] * expected + result = [] + for i in range(expected): + val = parsed[i] if i < len(parsed) else None + result.append(int(val) if isinstance(val, (int, float)) and val == int(val) else None) + return result + + +class L3LLMProcessor: + """LLM-based entity disambiguation — drop-in replacement for L3Processor.""" + + def __init__(self, config: L3LLMConfig): + self.config = config + self.schema = {} + self._client = self._build_client() + + def _build_client(self): + if self.config.provider == "gemini": + try: + import google.generativeai as genai + except ImportError: + raise ImportError("pip install google-generativeai") + api_key = self.config.api_key or os.environ.get("GEMINI_API_KEY") + if not api_key: + raise ValueError("Set GEMINI_API_KEY env var or config.api_key") + genai.configure(api_key=api_key) + return genai.GenerativeModel( + self.config.model_name, + system_instruction=_SYSTEM_PROMPT, + ) + elif self.config.provider == "openai": + try: + from openai import OpenAI + except ImportError: + raise ImportError("pip install openai") + api_key = self.config.api_key or os.environ.get("OPENAI_API_KEY") + return OpenAI(api_key=api_key) + else: + raise ValueError(f"Unknown provider: {self.config.provider!r}") + + def _complete(self, prompt: str) -> str: + if self.config.provider == "gemini": + response = self._client.generate_content( + prompt, + generation_config={"temperature": self.config.temperature}, + ) + return response.text + else: + response = self._client.chat.completions.create( + model=self.config.model_name, + messages=[ + {"role": "system", "content": _SYSTEM_PROMPT}, + {"role": "user", "content": prompt}, + ], + temperature=self.config.temperature, + ) + return response.choices[0].message.content + + def _call_with_retry(self, batch_items: list[dict]) -> list[int | None]: + prompt = _build_prompt(batch_items) + for attempt in range(self.config.max_retries): + try: + raw = self._complete(prompt) + return _parse_llm_response(raw, len(batch_items)) + except Exception as e: + logger.warning(f"LLM attempt {attempt + 1} failed: {e}") + if attempt < self.config.max_retries - 1: + time.sleep(self.config.retry_delay) + return [None] * len(batch_items) + + def _cand_info(self, c: Any) -> dict: + def get(attr: str) -> Any: + return getattr(c, attr, c.get(attr, "") if isinstance(c, dict) else "") + + return { + "label": get("label"), + "description": get("description"), + "entity_id": get("entity_id"), + } + + def __call__( + self, + texts: List[str], + candidates: List[List[List[Any]]], + l1_entities: List[List[Any]] | None = None, + input_data: Any = None, + ) -> L3Output: + all_entities: List[List[L3Entity]] = [] + + for text_idx, (text, per_span_candidates) in enumerate(zip(texts, candidates)): + l1 = l1_entities[text_idx] if l1_entities and text_idx < len(l1_entities) else [] + + items = [] + for entity, cands in zip(l1, per_span_candidates): + if not cands: + items.append(None) + continue + start = entity["start"] if isinstance(entity, dict) else entity.start + end = entity["end"] if isinstance(entity, dict) else entity.end + top_cands = cands[: self.config.top_k] + items.append( + { + "start": start, + "end": end, + "mention_text": text[start:end], + "context": _mark_mention(text, start, end), + "candidates": [self._cand_info(c) for c in top_cands], + "raw_candidates": top_cands, + } + ) + + non_null = [(i, item) for i, item in enumerate(items) if item is not None] + choices: dict[int, int | None] = {} + + for batch_start in range(0, len(non_null), self.config.batch_size): + batch = non_null[batch_start : batch_start + self.config.batch_size] + results = self._call_with_retry([item for _, item in batch]) + for local_idx, choice in enumerate(results): + choices[batch[local_idx][0]] = choice + + text_entities: List[L3Entity] = [] + for i, item in enumerate(items): + if item is None: + continue + choice = choices.get(i) + if choice is None or not (0 <= choice < len(item["raw_candidates"])): + continue + cand = item["raw_candidates"][choice] + label = getattr( + cand, "label", cand.get("label", "") if isinstance(cand, dict) else "" + ) + text_entities.append( + L3Entity( + text=item["mention_text"], + label=label, + start=item["start"], + end=item["end"], + score=1.0, + ) + ) + + all_entities.append(text_entities) + + return L3Output(entities=all_entities) + + +@processor_registry.register("l3_llm") +def create_l3_llm_processor(config_dict: dict, _pipeline: list | None = None) -> L3LLMProcessor: + config = L3LLMConfig(**config_dict) + return L3LLMProcessor(config) diff --git a/src/glinker/l4/processor.py b/src/glinker/l4/processor.py index 714faf3..fbdcad7 100644 --- a/src/glinker/l4/processor.py +++ b/src/glinker/l4/processor.py @@ -40,6 +40,7 @@ def __call__( texts: List[str] | None = None, candidates: List[List[Any]] | None = None, l1_entities: List[List[Any]] | None = None, + l3_entities: List[List[Any]] | None = None, input_data: L3Input = None, ) -> L3Output: """Process texts with candidate labels using chunked GLiNER inference. @@ -48,6 +49,8 @@ def __call__( 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 + l3_entities: Optional L3 entities per text - if provided, takes top + max_labels candidates sorted by L3 score instead of all L2 candidates input_data: Alternative L3Input object """ if texts is not None and candidates is not None: @@ -62,6 +65,40 @@ def __call__( all_entities = [] max_labels = self.config.max_labels + # If l3_entities provided, rerank only top L3 entities (not all L2 candidates) + if l3_entities is not None: + top_l3 = self._get_top_l3_entities(l3_entities, max_labels) + + for _idx, (text, text_top_l3) in enumerate(zip(texts_to_process, top_l3)): + if not text_top_l3: + all_entities.append([]) + continue + + # Build input_spans from L3 entities (specific spans to rerank) + l3_spans = {(e.start, e.end) for e in text_top_l3} + input_spans = [[{"start": e.start, "end": e.end} for e in text_top_l3]] + + # Get unique labels from L3 entities + labels = list(dict.fromkeys(e.label for e in text_top_l3)) + + # Run prediction on these spans + entities = self.component.predict_entities_chunked( + text, labels, max_labels, input_spans=input_spans + ) + + # Filter to only keep entities matching L3 spans + entities = [e for e in entities if (e.start, e.end) in l3_spans] + + # Apply remaining pipeline steps + 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) + + # Original behavior when no l3_entities provided # 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:] @@ -104,6 +141,25 @@ def __call__( return L3Output(entities=all_entities) + def _get_top_l3_entities( + self, + l3_entities: List[List[Any]], + max_labels: int, + ) -> List[List[Any]]: + """Get top L3 entities sorted by score. + + Returns top max_labels entities per text for reranking. + """ + result = [] + for text_l3_entities in l3_entities: + if not text_l3_entities: + result.append([]) + continue + # Sort by score descending and take top N + sorted_entities = sorted(text_l3_entities, key=lambda e: e.score, reverse=True) + result.append(sorted_entities[:max_labels]) + return result + def _extract_label(self, candidate: Any) -> str: """Extract label from candidate.""" if hasattr(candidate, "label"): diff --git a/tests/l3/test_processor.py b/tests/l3/test_processor.py index e67ffae..3f25eac 100644 --- a/tests/l3/test_processor.py +++ b/tests/l3/test_processor.py @@ -8,6 +8,7 @@ def create_processor(processor_name, config_dict): """Helper to create processor from registry.""" from glinker.core.registry import processor_registry + factory = processor_registry.get(processor_name) return factory(config_dict=config_dict, pipeline=None) @@ -21,16 +22,16 @@ def test_create_via_registry(self, l3_config_dict): def test_processor_has_component(self, l3_config_dict): processor = create_processor("l3_batch", l3_config_dict) - assert hasattr(processor, 'component') + assert hasattr(processor, "component") assert processor.component is not None def test_processor_has_config(self, l3_config_dict): processor = create_processor("l3_batch", l3_config_dict) - assert hasattr(processor, 'config') + assert hasattr(processor, "config") def test_processor_has_schema(self, l3_config_dict): processor = create_processor("l3_batch", l3_config_dict) - assert hasattr(processor, 'schema') + assert hasattr(processor, "schema") class TestL3ProcessorCall: @@ -49,8 +50,7 @@ def test_call_with_texts_and_candidates(self, l3_config_dict): ] result = processor( - texts=["TP53 causes cancer."], - candidates=[[candidates[0], candidates[1]]] + texts=["TP53 causes cancer."], candidates=[[candidates[0], candidates[1]]] ) assert isinstance(result, L3Output) @@ -84,12 +84,12 @@ def test_extract_label_from_record(self, l3_config_dict): processor = create_processor("l3_batch", l3_config_dict) record = DatabaseRecord(entity_id="1", label="TP53", description="A gene") - label = processor._extract_label(record) + label = processor._format_label(record, "{label}") assert label == "TP53" def test_extract_label_from_string(self, l3_config_dict): processor = create_processor("l3_batch", l3_config_dict) - label = processor._extract_label("simple_label") + label = processor._format_label("simple_label", "{label}") assert label == "simple_label" def test_create_gliner_labels_with_template(self, l3_config_dict): @@ -103,14 +103,15 @@ def test_create_gliner_labels_with_template(self, l3_config_dict): DatabaseRecord(entity_id="2", label="BRCA1", description="Breast cancer gene"), ] - labels, mapping = processor._create_gliner_labels_with_mapping(candidates) + labels, span_label_indices, mapping = processor._build_sparse_inputs([candidates]) assert len(labels) == 2 assert "TP53: Tumor protein" in labels assert "BRCA1: Breast cancer gene" in labels + assert span_label_indices == [[0, 1]] assert len(mapping) == 2 - def test_create_gliner_labels_deduplication(self, l3_config_dict): + def test_build_sparse_inputs_deduplicates_exact_labels(self, l3_config_dict): from glinker.l2.models import DatabaseRecord processor = create_processor("l3_batch", l3_config_dict) @@ -119,12 +120,15 @@ def test_create_gliner_labels_deduplication(self, l3_config_dict): candidates = [ DatabaseRecord(entity_id="1", label="TP53"), DatabaseRecord(entity_id="2", label="TP53"), # Duplicate - DatabaseRecord(entity_id="3", label="tp53"), # Case-insensitive duplicate + DatabaseRecord(entity_id="3", label="tp53"), ] - labels, mapping = processor._create_gliner_labels_with_mapping(candidates) + labels, span_label_indices, _mapping = processor._build_sparse_inputs([candidates]) - assert len(labels) == 1 # Should deduplicate + assert ( + len(labels) == 2 + ) # Exact duplicates are merged; differently-cased labels are distinct. + assert span_label_indices == [[0, 0, 1]] class TestL3ProcessorRanking: @@ -139,8 +143,8 @@ def test_rank_entities(self, l3_config_dict): "template": "{label}", "ranking": [ {"field": "gliner_score", "weight": 0.7}, - {"field": "popularity", "weight": 0.3} - ] + {"field": "popularity", "weight": 0.3}, + ], } entities = [ @@ -184,7 +188,7 @@ def test_can_use_precomputed_with_embeddings(self, l3_config_dict): entity_id="1", label="TP53", embedding=[0.1] * 768, - embedding_model_id=l3_config_dict["model_name"] + embedding_model_id=l3_config_dict["model_name"], ), ] @@ -201,7 +205,7 @@ def test_can_use_precomputed_model_mismatch(self, l3_config_dict): entity_id="1", label="TP53", embedding=[0.1] * 768, - embedding_model_id="different-model" # Wrong model + embedding_model_id="different-model", # Wrong model ), ]