feat: activate Focal Loss, add Dice Loss + span-width weighting, OpenVINO INT8 pipeline

gliner/model.py

@@ -8,7 +8,10 @@
 from pathlib import Path
 
 import torch
-import onnxruntime as ort
+try:
+    import onnxruntime as ort
+except (ImportError, Exception):
+    ort = None
 import transformers
 from tqdm import tqdm
 from torch import nn
@@ -51,7 +54,12 @@
     SpanGenerativeDecoder,
     TokenGenerativeDecoder,
 )
-from .training import Trainer, TrainingArguments
+# Lazy import — transformers.Trainer triggers torch.distributed import which
+# deadlocks on macOS ARM when multiple OpenMP runtimes are present.
+# Only load when train_model() is actually called.
+def _get_trainer_classes():
+    from .training import Trainer, TrainingArguments
+    return Trainer, TrainingArguments
 from .evaluation import BaseNEREvaluator, BaseRelexEvaluator
 from .onnx.model import (
     BaseORTModel,
@@ -101,10 +109,12 @@
 
 if is_module_available("onnxruntime"):
     import onnxruntime as ort
-
     ONNX_AVAILABLE = True
 else:
-    ONNX_AVAILABLE = False
+    ort = None
+    # export_to_onnx only needs torch.onnx.export (always available).
+    # onnxruntime is only required for ORT inference, not for the export itself.
+    ONNX_AVAILABLE = True
 
 logger = logging.getLogger(__name__)
 
@@ -1656,41 +1666,9 @@ def create_training_args(
         dataloader_num_workers: int = 1,
         report_to: str = "none",
         **kwargs,
-    ) -> TrainingArguments:
-        """Create training arguments with sensible defaults.
-
-        Args:
-            output_dir: Directory to save model checkpoints.
-            learning_rate: Learning rate for main parameters.
-            weight_decay: Weight decay for main parameters.
-            others_lr: Learning rate for other parameters.
-            others_weight_decay: Weight decay for other parameters.
-            focal_loss_alpha: Alpha for focal loss.
-            focal_loss_gamma: Gamma for focal loss.
-            rel_focal_loss_alpha: Alpha for relation focal loss. Defaults to entity alpha.
-            rel_focal_loss_gamma: Gamma for relation focal loss. Defaults to entity gamma.
-            focal_loss_prob_margin: Probability margin for focal loss.
-            loss_reduction: Loss reduction method.
-            negatives: Negative sampling ratio.
-            masking: Masking strategy.
-            lr_scheduler_type: Learning rate scheduler type.
-            warmup_ratio: Warmup ratio.
-            per_device_train_batch_size: Training batch size.
-            per_device_eval_batch_size: Evaluation batch size.
-            max_grad_norm: Maximum gradient norm.
-            max_steps: Maximum training steps.
-            save_steps: Save checkpoint every N steps.
-            save_total_limit: Maximum number of checkpoints to keep.
-            logging_steps: Log every N steps.
-            use_cpu: Whether to use CPU.
-            bf16: Whether to use bfloat16.
-            dataloader_num_workers: Number of dataloader workers.
-            report_to: Where to report metrics.
-            **kwargs: Additional training arguments.
-
-        Returns:
-            TrainingArguments instance.
-        """
+    ):
+        """Create training arguments with sensible defaults."""
+        _, TrainingArguments = _get_trainer_classes()
         return TrainingArguments(
             output_dir=output_dir,
             learning_rate=learning_rate,
@@ -1725,12 +1703,12 @@ def train_model(
         self,
         train_dataset,
         eval_dataset,
-        training_args: Optional[TrainingArguments] = None,
+        training_args=None,
         freeze_components: Optional[list[str]] = None,
         compile_model: bool = False,
         output_dir: Optional[Union[str, Path]] = None,
         **training_kwargs,
-    ) -> Trainer:
+    ):
         """Train the model.
 
         Args:
@@ -1778,6 +1756,7 @@ def train_model(
         else:
             trainer_kwargs["processing_class"] = self.data_processor.transformer_tokenizer
 
+        Trainer, _ = _get_trainer_classes()
         trainer = Trainer(**trainer_kwargs)
 
         # Train

gliner/modeling/base.py

@@ -39,7 +39,7 @@
 from .outputs import GLiNERBaseOutput, GLiNERRelexOutput, GLiNERDecoderOutput
 from .scorers import Scorer
 from .span_rep import SpanRepLayer
-from .loss_functions import cross_entropy_loss, focal_loss_with_logits
+from .loss_functions import cross_entropy_loss, focal_loss_with_logits, span_dice_loss
 from .multitask.triples_layers import TriplesScoreLayer
 from .multitask.relations_layers import RelationsRepLayer
 
@@ -183,6 +183,8 @@ def _loss(
         negatives: float = 1.0,
         masking: str = "none",
         normalize_prob: bool = True,
+        loss_type: str = "focal",
+        dice_gamma: float = 1.0,
     ) -> torch.Tensor:
         """Compute loss with optional negative sampling and masking.
 
@@ -203,16 +205,19 @@ def _loss(
         Returns:
             Loss tensor of same shape as labels.
         """
-        # Compute the loss per element using the focal loss function
-        all_losses = focal_loss_with_logits(
-            logits,
-            labels,
-            alpha=alpha,
-            gamma=gamma,
-            prob_margin=prob_margin,
-            label_smoothing=label_smoothing,
-            normalize_prob=normalize_prob,
-        )
+        if loss_type == "dice":
+            all_losses = span_dice_loss(logits, targets=labels, gamma=dice_gamma, reduction="none")
+        else:
+            # "focal" (default) and "bce" (alpha=-1, gamma=0) both go through focal_loss_with_logits
+            all_losses = focal_loss_with_logits(
+                logits,
+                labels,
+                alpha=alpha,
+                gamma=gamma,
+                prob_margin=prob_margin,
+                label_smoothing=label_smoothing,
+                normalize_prob=normalize_prob,
+            )
 
         # Create a mask of the same shape as labels:
         # For elements where labels==0, sample a Bernoulli random variable that is 1 with probability `negatives`
@@ -500,6 +505,9 @@ def loss(
         reduction: str = "sum",
         negatives: float = 1.0,
         masking: str = "none",
+        loss_type: str = "focal",
+        use_span_width_weight: bool = False,
+        dice_gamma: float = 1.0,
         **kwargs: Any,
     ) -> torch.Tensor:
         """Compute span classification loss.
@@ -516,6 +524,11 @@ def loss(
             reduction: Loss reduction method ('sum' or 'mean').
             negatives: Negative sampling probability.
             masking: Masking strategy for negative sampling.
+            loss_type: Loss function — 'focal' (default), 'bce', or 'dice'.
+            use_span_width_weight: If True, scale positive-span losses by
+                w(k) = 1 + log(k+1) where k is the 1-indexed span width.
+                Zero inference overhead; applied before final reduction.
+            dice_gamma: Self-adjustment exponent for Dice loss. Ignored for focal/bce.
             **kwargs: Additional arguments.
 
         Returns:
@@ -524,16 +537,40 @@ def loss(
         batch_size = scores.shape[0]
         num_classes = prompts_embedding_mask.shape[-1]
 
-        # Reshape scores and labels to match the expected shape
-        BS, _, _, CL = scores.shape
+        BS, L, K, CL = scores.shape
+
+        # Span-width-aware positive weighting: w(k) = 1 + log(k+1), k in [1..K].
+        # Applied before flattening so the K axis is still accessible.
+        if use_span_width_weight:
+            k_idx = torch.arange(1, K + 1, dtype=scores.dtype, device=scores.device)
+            width_weight = 1.0 + torch.log(k_idx)           # (K,)
+            # Broadcast to (B, L, K, C); boost positives only, leave negatives at 1.0
+            w = width_weight.view(1, 1, K, 1)
+            label_clamped = labels.clamp(min=0.0)
+            # weight tensor: 1 everywhere, boosted by w where label=1
+            pos_boost = 1.0 + (w - 1.0) * label_clamped    # (B, L, K, C)
+        else:
+            pos_boost = None
 
         scores = scores.view(BS, -1, CL)
         labels = labels.view(BS, -1, CL)
 
         all_losses = self._loss(
-            scores, labels, alpha, gamma, prob_margin, label_smoothing, negatives=negatives, masking=masking
+            scores,
+            labels,
+            alpha,
+            gamma,
+            prob_margin,
+            label_smoothing,
+            negatives=negatives,
+            masking=masking,
+            loss_type=loss_type,
+            dice_gamma=dice_gamma,
         )
 
+        if pos_boost is not None:
+            all_losses = all_losses * pos_boost.view(BS, -1, CL)
+
         masked_loss = all_losses.view(batch_size, -1, num_classes) * prompts_embedding_mask.unsqueeze(1)
         all_losses = masked_loss.view(-1, num_classes)
 

gliner/modeling/encoder.py

@@ -174,6 +174,7 @@ def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:
         """
         pair_attention_mask = kwargs.pop("pair_attention_mask", None)
         base_attention_mask = kwargs.pop("attention_mask", None)
+        kwargs.pop("token_lengths", None)  # GLiNER-internal kwarg, not accepted by HF models
         # Extract input_ids if present
         args = list(args)
         input_ids = kwargs.pop("input_ids", None)

gliner/modeling/loss_functions.py

@@ -2,6 +2,70 @@
 import torch.nn.functional as F
 
 
+def span_dice_loss(
+    inputs: torch.Tensor,
+    targets: torch.Tensor,
+    gamma: float = 1.0,
+    reduction: str = "none",
+    ignore_index: int = -100,
+    eps: float = 1e-6,
+) -> torch.Tensor:
+    """Per-element self-adjusting Dice loss for span-level NER.
+
+    Adapted from Li et al. (ACL 2020) "Dice Loss for Data-imbalanced NLP Tasks"
+    to operate element-wise over the (B, L*K, C) span-prediction tensor.
+
+    Unlike BCE, Dice is directly optimizing a surrogate of the F1 coefficient
+    and is immune to O-label domination because false negatives and false positives
+    are balanced in the denominator regardless of class prevalence.
+
+    The self-adjusting modulator (1-p)^gamma down-weights easy-to-classify
+    negative spans (high confidence, low p), focusing gradient on hard positives.
+
+    Loss per element:
+        adjusted_p = (1 - sigmoid(input))^gamma * sigmoid(input)
+        DSC_adj    = (2 * adjusted_p * y + eps) / (adjusted_p + y + eps)
+        loss       = 1 - DSC_adj
+
+    Args:
+        inputs: Predicted logits of arbitrary shape.
+        targets: Ground truth binary labels, same shape as inputs.
+            Values in {0, 1} or ignore_index.
+        gamma: Self-adjustment exponent. Higher values focus more on hard negatives.
+            Defaults to 1.0. Set to 0.0 to recover standard soft Dice.
+        reduction: 'none', 'mean', or 'sum'. Defaults to 'none'.
+        ignore_index: Target values to mask out from the loss. Defaults to -100.
+        eps: Numerical stability epsilon. Defaults to 1e-6.
+
+    Returns:
+        Loss tensor. Shape depends on reduction.
+    """
+    valid_mask = targets != ignore_index
+
+    p = torch.sigmoid(inputs)
+    y = targets.clamp(min=0.0)
+
+    adjusted_p = (1.0 - p).pow(gamma) * p
+
+    numerator = 2.0 * adjusted_p * y
+    denominator = adjusted_p + y
+
+    per_element = 1.0 - (numerator + eps) / (denominator + eps)
+    per_element = per_element * valid_mask.float()
+
+    if reduction == "none":
+        return per_element
+    elif reduction == "mean":
+        return per_element.sum() / valid_mask.float().sum().clamp(min=1)
+    elif reduction == "sum":
+        return per_element.sum()
+    else:
+        raise ValueError(
+            f"Invalid value for argument 'reduction': '{reduction}'. "
+            "Supported modes: 'none', 'mean', 'sum'"
+        )
+
+
 def focal_loss_with_logits(
     inputs: torch.Tensor,
     targets: torch.Tensor,

gliner/onnx/model.py

@@ -5,14 +5,18 @@
 span-level and token-level named entity recognition, as well as relation
 extraction models.
 """
+from __future__ import annotations
 
 import warnings
 from abc import ABC, abstractmethod
 from typing import Any, Dict, Optional
 
 import numpy as np
 import torch
-import onnxruntime as ort
+try:
+    import onnxruntime as ort
+except (ImportError, Exception):
+    ort = None
 
 from ..modeling.outputs import GLiNERBaseOutput, GLiNERRelexOutput
 

gliner/training/trainer.py

@@ -14,16 +14,14 @@
 import torch
 import transformers
 from torch import nn
-from transformers.trainer import (
-    get_parameter_names,
-    is_sagemaker_mp_enabled,
-)
 from transformers.trainer_utils import set_seed
-
-if is_sagemaker_mp_enabled():
-    from transformers.trainer_pt_utils import smp_forward_backward
 from torch.utils.data import Dataset, DataLoader
 
+
+def _get_trainer_imports():
+    from transformers.trainer import get_parameter_names, is_sagemaker_mp_enabled
+    return get_parameter_names, is_sagemaker_mp_enabled
+
 ALL_LAYERNORM_LAYERS = [nn.LayerNorm]
 
 logger = logging.getLogger(__name__)
@@ -85,6 +83,9 @@ class TrainingArguments(transformers.TrainingArguments):
     loss_reduction: Optional[str] = "sum"
     negatives: Optional[float] = 1.0
     masking: Optional[str] = "global"
+    loss_type: Optional[str] = "focal"
+    use_span_width_weight: Optional[bool] = False
+    dice_gamma: Optional[float] = 1.0
 
 
 class Trainer(transformers.Trainer):
@@ -164,6 +165,9 @@ def compute_loss(
             reduction=self.args.loss_reduction,
             negatives=self.args.negatives,
             masking=self.args.masking,
+            loss_type=self.args.loss_type,
+            use_span_width_weight=self.args.use_span_width_weight,
+            dice_gamma=self.args.dice_gamma,
             **inputs,
         )
 
@@ -184,7 +188,9 @@ def training_step(
             raise KeyError(f"Batch has no 'labels'. Keys: {list(inputs.keys())}")
 
         try:
+            _, is_sagemaker_mp_enabled = _get_trainer_imports()
             if is_sagemaker_mp_enabled():
+                from transformers.trainer_pt_utils import smp_forward_backward
                 loss_mb = smp_forward_backward(model, inputs, self.args.gradient_accumulation_steps)
                 return loss_mb.reduce_mean().detach().to(self.args.device)
 
@@ -225,6 +231,7 @@ def training_step(
             raise
 
     def create_optimizer(self):
+        get_parameter_names, is_sagemaker_mp_enabled = _get_trainer_imports()
         if is_sagemaker_mp_enabled():
             return super().create_optimizer()