Merge pull request #326 from urchade/fix/training

Fix training issues and compatability with transformers v4.x and v5

Author: Ingvarstep

gliner/config.py

@@ -324,7 +324,7 @@ def __init__(
     def model_type(self):
         """Auto-detect model type based on configuration."""
         if self.labels_decoder:
-            if self.span_mode == 'token-level':
+            if self.span_mode == "token-level":
                 return "gliner_uni_encoder_token_decoder"
             else:
                 return "gliner_uni_encoder_span_decoder"
@@ -357,4 +357,4 @@ def model_type(self):
         "gliner_bi_encoder_span": BiEncoderSpanConfig,
         "gliner_bi_encoder_token": BiEncoderTokenConfig,
     }
-)
+)

gliner/data_processing/processor.py

@@ -807,6 +807,8 @@ def create_labels(self, batch):
 
         for i, sentence_entities in enumerate(batch["entities"]):
             for st, ed, sp_label in sentence_entities:
+                if sp_label not in batch["classes_to_id"][i]:
+                    continue
                 lbl = batch["classes_to_id"][i][sp_label]
                 class_idx = lbl - 1  # Convert to 0-indexed
 

gliner/model.py

@@ -9,8 +9,10 @@
 
 import torch
 import onnxruntime as ort
+import transformers
 from tqdm import tqdm
 from torch import nn
+from packaging import version
 from safetensors import safe_open
 from transformers import AutoTokenizer
 from huggingface_hub import PyTorchModelHubMixin, snapshot_download
@@ -338,6 +340,14 @@ def _load_config(cls, config_file: Path, **config_overrides) -> object:
 
         return config
 
+    @staticmethod
+    def _set_tokenizer_spec_tokens(tokenizer):
+        if hasattr(tokenizer, "add_bos_token"):
+            tokenizer.add_bos_token = tokenizer.bos_token_id is not None
+        if hasattr(tokenizer, "add_eos_token"):
+            tokenizer.add_eos_token = tokenizer.eos_token_id is not None
+        return tokenizer
+
     @classmethod
     def _load_tokenizer(cls, config: GLiNERConfig, model_dir: Path, cache_dir: Optional[Path] = None):
         """
@@ -351,11 +361,14 @@ def _load_tokenizer(cls, config: GLiNERConfig, model_dir: Path, cache_dir: Optio
         Returns:
             Tokenizer instance or None
         """
-        if os.path.exists(model_dir / "tokenizer_config.json"):
-            return AutoTokenizer.from_pretrained(model_dir, cache_dir=cache_dir)
+        tokenizer_config_path = model_dir / "tokenizer_config.json"
+
+        if tokenizer_config_path.is_file():
+            tokenizer = AutoTokenizer.from_pretrained(model_dir, cache_dir=cache_dir)
         else:
-            return AutoTokenizer.from_pretrained(config.model_name, cache_dir=cache_dir)
-        return None
+            tokenizer = AutoTokenizer.from_pretrained(config.model_name, cache_dir=cache_dir)
+
+        return cls._set_tokenizer_spec_tokens(tokenizer)
 
     @classmethod
     def _load_state_dict(cls, model_file: Path, map_location: str = "cpu"):
@@ -514,7 +527,7 @@ def load_from_config(
         tokenizer = None
         if load_tokenizer:
             tokenizer = AutoTokenizer.from_pretrained(config_instance.model_name, cache_dir=cache_dir)
-
+            cls._set_tokenizer_spec_tokens(tokenizer)
         # Create model instance from scratch
         instance = cls(
             config_instance,
@@ -1110,15 +1123,22 @@ def train_model(
         # Create data collator
         data_collator = self._create_data_collator()
 
-        # Create trainer
-        trainer = Trainer(
-            model=self,
-            args=training_args,
-            train_dataset=train_dataset,
-            eval_dataset=eval_dataset,
-            tokenizer=self.data_processor.transformer_tokenizer,
-            data_collator=data_collator,
-        )
+        # Create trainer with version-conditional tokenizer argument
+        # transformers < 5.0 requires tokenizer, >= 5.0 does not
+        trainer_kwargs = {
+            "model": self,
+            "args": training_args,
+            "train_dataset": train_dataset,
+            "eval_dataset": eval_dataset,
+            "data_collator": data_collator,
+        }
+
+        if version.parse(transformers.__version__) < version.parse("5.0.0"):
+            trainer_kwargs["tokenizer"] = self.data_processor.transformer_tokenizer
+        else:
+            trainer_kwargs["processing_class"] = self.data_processor.transformer_tokenizer
+
+        trainer = Trainer(**trainer_kwargs)
 
         # Train
         trainer.train()
@@ -1134,6 +1154,7 @@ def _create_model(self, config, backbone_from_pretrained, cache_dir, **kwargs):
     def _create_data_processor(self, config, cache_dir, tokenizer=None, words_splitter=None, **kwargs):
         if tokenizer is None:
             tokenizer = AutoTokenizer.from_pretrained(config.model_name, cache_dir=cache_dir)
+            self._set_tokenizer_spec_tokens(tokenizer)
         self.data_processor = self.data_processor_class(config, tokenizer, words_splitter)
         return self.data_processor
 
@@ -1520,6 +1541,7 @@ def _create_data_processor(self, config, cache_dir, tokenizer=None, words_splitt
         labels_tokenizer = AutoTokenizer.from_pretrained(config.labels_encoder, cache_dir=cache_dir)
         if tokenizer is None:
             tokenizer = AutoTokenizer.from_pretrained(config.model_name, cache_dir=cache_dir)
+            self._set_tokenizer_spec_tokens(tokenizer)
 
         self.data_processor = self.data_processor_class(
             config, tokenizer, words_splitter, labels_tokenizer=labels_tokenizer
@@ -1981,6 +2003,7 @@ def _create_data_processor(self, config, cache_dir, tokenizer=None, words_splitt
         """Create data processor with decoder tokenizer."""
         if tokenizer is None:
             tokenizer = AutoTokenizer.from_pretrained(config.model_name, cache_dir=cache_dir)
+            self._set_tokenizer_spec_tokens(tokenizer)
 
         if words_splitter is None:
             words_splitter = WordsSplitter(config.words_splitter_type)
@@ -2242,6 +2265,7 @@ def _create_data_processor(self, config, cache_dir, tokenizer=None, words_splitt
         """Create relation extraction data processor."""
         if tokenizer is None:
             tokenizer = AutoTokenizer.from_pretrained(config.model_name, cache_dir=cache_dir)
+            self._set_tokenizer_spec_tokens(tokenizer)
 
         if words_splitter is None:
             words_splitter = WordsSplitter(config.words_splitter_type)
@@ -2271,7 +2295,7 @@ def inference(
         self,
         texts: Union[str, List[str]],
         labels: List[str],
-        relations: List[str],
+        relations: List[str] = [],
         flat_ner: bool = True,
         threshold: float = 0.5,
         adjacency_threshold: Optional[float] = None,