removed protobuf + transformers 5.x

Author: jogonba2

.github/workflows/ci.yaml

@@ -8,7 +8,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [3.8, 3.9]
+        python-version: [3.8, 3.9, 3.10]
 
     steps:
       - uses: actions/checkout@v2

comet/encoders/base.py

@@ -80,6 +80,18 @@ def from_pretrained(cls, pretrained_model: str):
         """
         raise NotImplementedError
 
+    @abc.abstractmethod
+    def build_inputs_with_special_tokens(self, token_ids_0: List[int], token_ids_1: List[int]) -> List[int]:
+        """Concatenate ids from two sequences.
+
+        This ensures compatibility with Transformers 4.x and 5.x, as the
+        `build_inputs_with_special_tokens method` has been removed from `PreTrainedTokenizerBase`.
+
+        Returns:
+            List[int]: an encoded sequence.
+        """
+        pass
+
     def freeze(self) -> None:
         """Frezees the entire encoder."""
         for param in self.parameters():
@@ -304,7 +316,7 @@ def concat_sequences(
                 torch.zeros(len(new_sequence[1:-1]) + 2, dtype=torch.int)
             )
             for j in range(1, len(inputs)):
-                new_sequence = self.tokenizer.build_inputs_with_special_tokens(
+                new_sequence = self.build_inputs_with_special_tokens(
                     new_sequence[1:-1], concat_input_ids[j][i][1:-1]
                 )
             if sum(lengths) > self.max_positions - special_tokens:

comet/encoders/bert.py

@@ -17,7 +17,7 @@
 ==============
     Pretrained BERT encoder from Hugging Face.
 """
-from typing import Dict, Optional
+from typing import Dict, List, Optional
 
 import torch
 from transformers import BertConfig, BertModel, BertTokenizerFast
@@ -115,6 +115,23 @@ def freeze_embeddings(self) -> None:
         for param in self.model.embeddings.parameters():
             param.requires_grad = False
 
+    def build_inputs_with_special_tokens(
+        self, token_ids_0: List[int], token_ids_1: List[int]
+    ) -> List[int]:
+        """Concatenate ids from two sequences.
+
+        Returns:
+            List[int]: an encoded sequence.
+        """
+        return (
+            [self.tokenizer.cls_token_id]
+            + token_ids_0
+            + [self.tokenizer.sep_token_id]
+            + token_ids_1
+            + [self.tokenizer.sep_token_id]
+        )
+
+
     def layerwise_lr(self, lr: float, decay: float):
         """Calculates the learning rate for each layer by applying a small decay.
 
@@ -168,15 +185,23 @@ def forward(
             Dict[str, torch.Tensor]: dictionary with 'sentemb', 'wordemb', 'all_layers'
                 and 'attention_mask'.
         """
-        last_hidden_states, pooler_output, all_layers = self.model(
+        output = self.model(
             input_ids=input_ids,
             token_type_ids=token_type_ids,
             attention_mask=attention_mask,
             output_hidden_states=True,
             return_dict=False,
         )
+
+        if len(output) == 2:
+            last_hidden_states, all_layers = output
+            sentemb = last_hidden_states[:, 0, :]
+        else:
+            last_hidden_states, pooler_output, all_layers = output
+            sentemb = pooler_output if pooler_output is not None else last_hidden_states[:, 0, :]
+
         return {
-            "sentemb": pooler_output,
+            "sentemb": sentemb,
             "wordemb": last_hidden_states,
             "all_layers": all_layers,
             "attention_mask": attention_mask,

comet/encoders/rembert.py

@@ -18,6 +18,8 @@
     Pretrained RemBERT encoder from Google. This encoder is similar to BERT but uses 
     sentencepiece like XLMR.
 """
+from typing import List
+
 from transformers import RemBertConfig, RemBertModel, RemBertTokenizerFast
 
 from comet.encoders.xlmr import Encoder, XLMREncoder
@@ -63,6 +65,22 @@ def size_separator(self):
     def uses_token_type_ids(self):
         return True
 
+    def build_inputs_with_special_tokens(
+        self, token_ids_0: List[int], token_ids_1: List[int]
+    ) -> List[int]:
+        """Concatenate ids from two sequences.
+
+        Returns:
+            List[int]: an encoded sequence.
+        """
+        return (
+            [self.tokenizer.cls_token_id]
+            + token_ids_0
+            + [self.tokenizer.sep_token_id]
+            + token_ids_1
+            + [self.tokenizer.sep_token_id]
+        )
+
     @classmethod
     def from_pretrained(
         cls,

comet/encoders/xlmr.py

@@ -17,7 +17,7 @@
 ==============
     Pretrained XLM-RoBERTa  encoder from Hugging Face.
 """
-from typing import Dict
+from typing import Dict, List
 
 import torch
 from transformers import XLMRobertaConfig, XLMRobertaModel, XLMRobertaTokenizerFast
@@ -69,6 +69,22 @@ def size_separator(self):
     def uses_token_type_ids(self):
         return False
 
+    def build_inputs_with_special_tokens(
+        self, token_ids_0: List[int], token_ids_1: List[int]
+    ) -> List[int]:
+        """Concatenate ids from two sequences.
+
+        Returns:
+            List[int]: an encoded sequence.
+        """
+        return (
+            [self.tokenizer.cls_token_id]
+            + token_ids_0
+            + [self.tokenizer.sep_token_id] * 2
+            + token_ids_1
+            + [self.tokenizer.sep_token_id]
+        )
+
     @classmethod
     def from_pretrained(
         cls,
@@ -92,12 +108,18 @@ def from_pretrained(
     def forward(
         self, input_ids: torch.Tensor, attention_mask: torch.Tensor, **kwargs
     ) -> Dict[str, torch.Tensor]:
-        last_hidden_states, _, all_layers = self.model(
+        output = self.model(
             input_ids=input_ids,
             attention_mask=attention_mask,
             output_hidden_states=True,
             return_dict=False,
         )
+
+        if len(output) == 2:
+            last_hidden_states, all_layers = output
+        else:
+            last_hidden_states, _, all_layers = output
+
         return {
             "sentemb": last_hidden_states[:, 0, :],
             "wordemb": last_hidden_states,

comet/models/predict_writer.py

@@ -97,14 +97,18 @@ def flatten_predictions(predictions):
         files = sorted(os.listdir(self.output_dir))
         pred = flatten_predictions(
             [
-                flatten_predictions(torch.load(os.path.join(self.output_dir, f)))
+                flatten_predictions(
+                    torch.load(os.path.join(self.output_dir, f), weights_only=False)
+                )
                 for f in files
                 if "pred" in f
             ]
         )
         indices = flatten(
             [
-                flatten(torch.load(os.path.join(self.output_dir, f))[0])
+                flatten(
+                    torch.load(os.path.join(self.output_dir, f), weights_only=False)[0]
+                )
                 for f in files
                 if "batch_indices" in f
             ]