Merge pull request #351 from maxwbuckley/fix-test-suite

Fix 35 pre-existing test failures; add CI workflow

Author: Ingvarstep

.github/workflows/tests.yml

@@ -0,0 +1,67 @@
+name: Tests
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+    branches:
+      - main
+  workflow_dispatch:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  test:
+    name: pytest (Python ${{ matrix.python-version }})
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.10", "3.12"]
+
+    steps:
+      - name: Check out repository
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Cache pip
+        uses: actions/cache@v4
+        with:
+          path: ~/.cache/pip
+          key: ${{ runner.os }}-py${{ matrix.python-version }}-pip-${{ hashFiles('requirements.txt', 'pyproject.toml') }}
+          restore-keys: |
+            ${{ runner.os }}-py${{ matrix.python-version }}-pip-
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install -r requirements.txt
+          pip install pytest pytest-asyncio sentencepiece onnxruntime
+
+      - name: Run pytest
+        run: pytest -q --tb=short
+
+  lint:
+    name: ruff
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check out repository
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+
+      - name: Install ruff
+        run: pip install ruff
+
+      - name: ruff check
+        run: ruff check gliner

gliner/config.py

@@ -230,6 +230,7 @@ def __init__(
                 the per-type entity drop probability. Defaults to (0.0, 0.4).
             augment_rel_drop_prob (tuple, optional): Range (min, max) from which to sample
                 the per-type relation drop probability. Defaults to (0.0, 0.4).
+            augment_add_other_prob (float, optional): Probability of adding "other" relation to a pair with no relation.
             rel_id_to_classes (Optional[dict]): Mapping from relation class IDs to class names. Defaults to None.
             **kwargs: Additional keyword arguments passed to UniEncoderConfig.
 

gliner/data_processing/processor.py

@@ -1368,7 +1368,8 @@ def create_labels(self, batch, blank=None):
             if self.config.decoder_mode == "span":
                 # Collect entity labels in order of appearance
                 sorted_entities = sorted(ner, key=lambda x: (x[0], x[1])) if ner else []
-                for start, end, label in sorted_entities:
+                # start, end, label = entity
+                for _, end, label in sorted_entities:
                     if label in classes_to_id and end < num_tokens:
                         decoder_label_strings.append(label)
             elif self.config.decoder_mode == "prompt":
@@ -1477,8 +1478,8 @@ def augment_example(self, example, ner_negatives=None, other_keyword="other"):
         rel_drop_prob = random.uniform(*self.config.augment_rel_drop_prob)
         add_other = random.random() < self.config.augment_add_other_prob
 
-        all_ent_types = set(e[-1] for e in ner)
-        all_rel_types = set(r[-1] for r in relations) if relations else set()
+        all_ent_types = {e[-1] for e in ner}
+        all_rel_types = {r[-1] for r in relations} if relations else set()
 
         # "other" is exempt from dropping since it's our replacement label
         dropped_ent_types = {t for t in all_ent_types if t != other_keyword and random.random() < ent_drop_prob}
@@ -1512,7 +1513,7 @@ def augment_example(self, example, ner_negatives=None, other_keyword="other"):
             old_to_new_idx[i] = len(new_ner)
             if ent_type in dropped_ent_types and add_other:
                 # Replace dropped type with "other"
-                new_ner.append(list(ent[:-1]) + [other_keyword])
+                new_ner.append([*ent[:-1], other_keyword])
             else:
                 new_ner.append(ent)
 
@@ -1649,13 +1650,12 @@ def collate_raw_batch(
             relation extraction.
         """
         # Apply data augmentation if enabled (only during dynamic mapping generation)
-        augment_prob = getattr(self.config, 'augment_data_prob', 0.0)
+        augment_prob = getattr(self.config, "augment_data_prob", 0.0)
         if augment_prob > 0.0 and class_to_ids is None and entity_types is None:
             if ner_negatives is None:
                 ner_negatives = get_negatives(batch_list, sampled_neg=100, key="ner")
             batch_list = [
-                self.augment_example(b, ner_negatives) if random.random() < augment_prob else b
-                for b in batch_list
+                self.augment_example(b, ner_negatives) if random.random() < augment_prob else b for b in batch_list
             ]
         if class_to_ids is None and entity_types is None:
             # Dynamically infer per-example mappings
@@ -1768,7 +1768,11 @@ def preprocess_example(self, tokens, ner, classes_to_id, relations, rel_classes_
                 head_idx, tail_idx, rel_type = rel
 
                 # Use compact indices so rel_idx aligns with target_span_rep positions
-                if head_idx in entity_to_compact_idx and tail_idx in entity_to_compact_idx and rel_type in rel_classes_to_id:
+                if (
+                    head_idx in entity_to_compact_idx
+                    and tail_idx in entity_to_compact_idx
+                    and rel_type in rel_classes_to_id
+                ):
                     rel_idx_list.append([entity_to_compact_idx[head_idx], entity_to_compact_idx[tail_idx]])
                     rel_label_list.append(rel_classes_to_id[rel_type])
 
@@ -1833,7 +1837,9 @@ def create_batch_dict(self, batch, class_to_ids, id_to_classes, rel_class_to_ids
             "rel_id_to_classes": rel_id_to_classes,
         }
 
-    def create_relation_labels(self, batch, add_reversed_negatives=True, add_random_negatives=True, negative_ratio=(1.0, 10.0)):
+    def create_relation_labels(
+        self, batch, add_reversed_negatives=True, add_random_negatives=True, negative_ratio=(1.0, 10.0)
+    ):
         """Create relation labels with negative pair sampling.
 
         Overrides the span-based version to work with token-level entity representations.
@@ -1870,7 +1876,7 @@ def create_relation_labels(self, batch, add_reversed_negatives=True, add_random_
 
         # Batch CPU transfer to avoid per-element .item() sync
         batch_ents_cpu = batch_ents.tolist()
-        max_En = max(max(batch_ents_cpu), 1)
+        max_En = max(*batch_ents_cpu, 1)
 
         rel_class_to_ids = batch["rel_class_to_ids"]
         if isinstance(rel_class_to_ids, list):

gliner/decoding/decoder.py

@@ -20,6 +20,7 @@ class Span:
         class_probs: Optional dict of top-k class probabilities
         generated_labels: Optional list of generated labels (for generative decoders)
     """
+
     start: int
     end: int
     entity_type: str
@@ -260,7 +261,7 @@ def _decode_batch_item(
         """
         # Mask probabilities to only include input spans (for efficiency)
         if input_spans_i is not None:
-            L, K_dim, C = probs_i.shape
+            L, K_dim, _ = probs_i.shape
             span_filter = torch.zeros(L, K_dim, dtype=torch.bool, device=probs_i.device)
             for word_start, word_end in input_spans_i:
                 width = word_end - word_start
@@ -358,18 +359,20 @@ class IDs to class names.
         if B == 1:
             id_to_class_0 = self._get_id_to_class_for_sample(id_to_classes, 0)
             input_spans_0 = input_spans[0] if input_spans is not None else None
-            return [self._decode_batch_item(
-                probs_i=probs[0],
-                tokens_i=tokens[0],
-                id_to_class_i=id_to_class_0,
-                K=K,
-                threshold=threshold,
-                flat_ner=flat_ner,
-                multi_label=multi_label,
-                span_label_map=span_label_maps[0],
-                return_class_probs=return_class_probs,
-                input_spans_i=input_spans_0,
-            )]
+            return [
+                self._decode_batch_item(
+                    probs_i=probs[0],
+                    tokens_i=tokens[0],
+                    id_to_class_i=id_to_class_0,
+                    K=K,
+                    threshold=threshold,
+                    flat_ner=flat_ner,
+                    multi_label=multi_label,
+                    span_label_map=span_label_maps[0],
+                    return_class_probs=return_class_probs,
+                    input_spans_i=input_spans_0,
+                )
+            ]
 
         # Apply input_spans mask at batch level (one mask, one multiply)
         if input_spans is not None:
@@ -392,9 +395,7 @@ class IDs to class names.
             return [[] for _ in range(B)]
 
         # ONE vectorized valid-span check across entire batch
-        num_tokens = torch.tensor(
-            [len(t) for t in tokens], device=probs.device, dtype=torch.long
-        )
+        num_tokens = torch.tensor([len(t) for t in tokens], device=probs.device, dtype=torch.long)
         valid = (s_idx + k_idx + 1) <= num_tokens[b_idx]
         b_idx = b_idx[valid]
         s_idx = s_idx[valid]
@@ -427,15 +428,11 @@ class IDs to class names.
             top_indices_list = all_top_indices.tolist()
 
         # Pre-resolve id_to_class mappings per batch item
-        id_to_class_per_item = [
-            self._get_id_to_class_for_sample(id_to_classes, i) for i in range(B)
-        ]
+        id_to_class_per_item = [self._get_id_to_class_for_sample(id_to_classes, i) for i in range(B)]
 
         # Group by batch item and build Span objects (pure Python)
         batch_spans: List[List[Span]] = [[] for _ in range(B)]
-        for j, (b, s, k, c, flat_idx, score) in enumerate(
-            zip(b_list, s_list, k_list, c_list, flat_idxs, scores)
-        ):
+        for j, (b, s, k, c, flat_idx, score) in enumerate(zip(b_list, s_list, k_list, c_list, flat_idxs, scores)):
             id_to_class_i = id_to_class_per_item[b]
 
             class_probs = None
@@ -445,16 +442,11 @@ class IDs to class names.
                     class_name = id_to_class_i.get(idx + 1, f"class_{idx}")
                     class_probs[class_name] = prob
 
-            span = self._build_span_tuple(
-                s, k, c, flat_idx, score, id_to_class_i, span_label_maps[b], class_probs
-            )
+            span = self._build_span_tuple(s, k, c, flat_idx, score, id_to_class_i, span_label_maps[b], class_probs)
             batch_spans[b].append(span)
 
         # Per-item greedy search (inherently sequential, but cheap pure Python)
-        return [
-            self.greedy_search(spans, flat_ner, multi_label=multi_label)
-            for spans in batch_spans
-        ]
+        return [self.greedy_search(spans, flat_ner, multi_label=multi_label) for spans in batch_spans]
 
     def decode(
         self,
@@ -544,13 +536,7 @@ def _build_span_tuple(
             Span: Span object with entity properties.
         """
         ent_type = id_to_class[class_idx + 1]  # +1 because 0 is <pad>
-        return Span(
-            start=start,
-            end=start + width,
-            entity_type=ent_type,
-            score=score,
-            class_probs=class_probs
-        )
+        return Span(start=start, end=start + width, entity_type=ent_type, score=score, class_probs=class_probs)
 
 
 class SpanGenerativeDecoder(BaseSpanDecoder):
@@ -679,7 +665,7 @@ def _build_span_tuple(
             entity_type=ent_type,
             score=score,
             class_probs=class_probs,
-            generated_labels=gen_ent_type
+            generated_labels=gen_ent_type,
         )
 
     def decode_generative(
@@ -864,15 +850,8 @@ def _decode_relations_batch(
     # 3. Vectorized index-validity check
     head = rel_idx[..., 0]  # (B, R)
     tail = rel_idx[..., 1]  # (B, R)
-    num_spans = torch.tensor(
-        [len(s) for s in spans], device=rel_idx.device, dtype=head.dtype
-    )  # (B,)
-    valid = (
-        (head >= 0)
-        & (tail >= 0)
-        & (head < num_spans[:, None])
-        & (tail < num_spans[:, None])
-    )  # (B, R)
+    num_spans = torch.tensor([len(s) for s in spans], device=rel_idx.device, dtype=head.dtype)  # (B,)
+    valid = (head >= 0) & (tail >= 0) & (head < num_spans[:, None]) & (tail < num_spans[:, None])  # (B, R)
     rel_probs = rel_probs * valid.unsqueeze(-1)
 
     # 4. Single torch.where on the full (B, R, C) tensor
@@ -898,9 +877,7 @@ def _decode_relations_batch(
         mapping = rel_id_to_classes[b] if is_list else rel_id_to_classes
         if c1 not in mapping:
             continue
-        relations[b].append(
-            (int(head_list[k]), mapping[c1], int(tail_list[k]), scores[k])
-        )
+        relations[b].append((int(head_list[k]), mapping[c1], int(tail_list[k]), scores[k]))
 
     return relations
 
@@ -955,13 +932,7 @@ def _build_span_tuple(
             Span: Span object with entity properties.
         """
         ent_type = id_to_class[class_idx + 1]  # +1 because 0 is <pad>
-        return Span(
-            start=start,
-            end=start + width,
-            entity_type=ent_type,
-            score=score,
-            class_probs=class_probs
-        )
+        return Span(start=start, end=start + width, entity_type=ent_type, score=score, class_probs=class_probs)
 
     def _build_entity_span_to_decoded_idx(
         self,
@@ -1151,6 +1122,7 @@ def decode(
             rel_idx: Optional tensor of shape (batch_size, num_relations, 2).
             rel_logits: Optional tensor of shape (batch_size, num_relations, num_relation_classes).
             rel_mask: Optional boolean tensor of shape (batch_size, num_relations).
+            return_class_probs: Whether to include class probabilities in the decoded spans.
             flat_ner: If True, applies greedy filtering for non-overlapping entities.
             threshold: Minimum confidence score for entity predictions.
             relation_threshold: Minimum confidence score for relation predictions.
@@ -1266,13 +1238,8 @@ def _calculate_span_score(
                     start_score = start_cpu[st][cls_st]
                     end_score = end_cpu[ed][cls_ed]
                     # The span score is the minimum value among all scores
-                    spn_score = min(min(ins), start_score, end_score)
-                    span_i.append(Span(
-                        start=st,
-                        end=ed,
-                        entity_type=id_to_classes[cls_st + 1],
-                        score=spn_score
-                    ))
+                    spn_score = min(*ins, start_score, end_score)
+                    span_i.append(Span(start=st, end=ed, entity_type=id_to_classes[cls_st + 1], score=spn_score))
         return span_i
 
     def _decode_from_spans(
@@ -1349,12 +1316,7 @@ class IDs to class names.
                     class_id = class_idx + 1  # Convert to 1-indexed
                     if class_id in id_to_class_i:
                         entity_type = id_to_class_i[class_id]
-                        span_scores.append(Span(
-                            start=span_start,
-                            end=span_end,
-                            entity_type=entity_type,
-                            score=prob
-                        ))
+                        span_scores.append(Span(start=span_start, end=span_end, entity_type=entity_type, score=prob))
 
             # Apply greedy search to handle overlapping spans if needed
             span_i = self.greedy_search(span_scores, flat_ner, multi_label)
@@ -1664,6 +1626,8 @@ def decode(
             rel_id_to_classes: Optional mapping from relation class IDs to relation names.
                 If None, relation decoding is skipped and empty relation lists are returned.
                 Can be either a single Dict or List[Dict] for per-sample mappings.
+            entity_spans: Optional tensor of pre-computed entity spans to use instead
+                of decoding them from model_output.
                 Class IDs are 1-indexed.
             **kwargs: Additional keyword arguments passed to the parent class decode method.