Upload tokenizer

tokenization_kpr.py

@@ -255,8 +255,7 @@ def preprocess_text(
 ) -> dict[str, list[int]]:
     tokens = []
     entity_ids = []
-    entity_start_positions = []
-    entity_lengths = []
+    entity_position_ids = []
     if title is not None:
         if title_mentions is None:
             title_mentions = []
@@ -265,8 +264,7 @@ def preprocess_text(
         tokens += title_tokens + [tokenizer.sep_token]
         for entity in title_entities:
             entity_ids.append(entity.entity_id)
-            entity_start_positions.append(entity.start)
-            entity_lengths.append(entity.end - entity.start)
+            entity_position_ids.append(list(range(entity.start, entity.end)))
 
     if mentions is None:
         mentions = []
@@ -276,16 +274,14 @@ def preprocess_text(
     tokens += text_tokens
     for entity in text_entities:
         entity_ids.append(entity.entity_id)
-        entity_start_positions.append(entity.start + entity_offset)
-        entity_lengths.append(entity.end - entity.start)
+        entity_position_ids.append(list(range(entity.start + entity_offset, entity.end + entity_offset)))
 
     input_ids = tokenizer.convert_tokens_to_ids(tokens)
 
     return {
         "input_ids": input_ids,
         "entity_ids": entity_ids,
-        "entity_start_positions": entity_start_positions,
-        "entity_lengths": entity_lengths,
+        "entity_position_ids": entity_position_ids,
     }
 
 
@@ -349,8 +345,7 @@ class KPRBertTokenizer(BertTokenizer):
         "token_type_ids",
         "attention_mask",
         "entity_ids",
-        "entity_start_positions",
-        "entity_lengths",
+        "entity_position_ids",
     ]
 
     def __init__(
@@ -379,7 +374,7 @@ class KPRBertTokenizer(BertTokenizer):
                     "Make sure `embeddings.py` and `entity_vocab.tsv` are consistent."
                 )
 
-    def _preprocess_text(self, text: str, **kwargs) -> dict[str, list[int]]:
+    def _preprocess_text(self, text: str, **kwargs) -> dict[str, list[int | list[int]]]:
         mentions = self.entity_linker.detect_mentions(text)
         model_inputs = preprocess_text(
             text=text,
@@ -395,18 +390,26 @@ class KPRBertTokenizer(BertTokenizer):
         # We exclude "return_tensors" from kwargs
         # to avoid issues in passing the data to BatchEncoding outside this method
         prepared_inputs = self.prepare_for_model(
-            model_inputs["input_ids"], **{k: v for k, v in kwargs.items() if k != "return_tensors"}
+            model_inputs["input_ids"],
+            **{k: v for k, v in kwargs.items() if k != "return_tensors"},
         )
         model_inputs.update(prepared_inputs)
 
         # Account for special tokens
-        if kwargs.get("add_special_tokens"):
+        if kwargs.get("add_special_tokens", True):
             if prepared_inputs["input_ids"][0] != self.cls_token_id:
                 raise ValueError(
                     "We assume that the input IDs start with the [CLS] token with add_special_tokens = True."
                 )
-            # Shift the entity start positions by 1 to account for the [CLS] token
-            model_inputs["entity_start_positions"] = [pos + 1 for pos in model_inputs["entity_start_positions"]]
+            # Shift the entity position IDs by 1 to account for the [CLS] token
+            model_inputs["entity_position_ids"] = [
+                [pos + 1 for pos in positions] for positions in model_inputs["entity_position_ids"]
+            ]
+
+        # If there is no entities in the text, we output padding entity for the model
+        if not model_inputs["entity_ids"]:
+            model_inputs["entity_ids"] = [0]  # The padding entity id is 0
+            model_inputs["entity_position_ids"] = [[0]]
 
         # Count the number of special tokens at the end of the input
         num_special_tokens_at_end = 0
@@ -414,26 +417,25 @@ class KPRBertTokenizer(BertTokenizer):
         if isinstance(input_ids, torch.Tensor):
             input_ids = input_ids.tolist()
         for input_id in input_ids[::-1]:
-            if int(input_id) not in {self.sep_token_id, self.pad_token_id, self.cls_token_id}:
+            if int(input_id) not in {
+                self.sep_token_id,
+                self.pad_token_id,
+                self.cls_token_id,
+            }:
                 break
             num_special_tokens_at_end += 1
 
         # Remove entities that are not in truncated input
         max_effective_pos = len(model_inputs["input_ids"]) - num_special_tokens_at_end
         entity_indices_to_keep = list()
-        for i, (start_pos, length) in enumerate(
-            zip(model_inputs["entity_start_positions"], model_inputs["entity_lengths"])
-        ):
-            if (start_pos + length) <= max_effective_pos:
+        for i, position_ids in enumerate(model_inputs["entity_position_ids"]):
+            if len(position_ids) > 0 and max(position_ids) < max_effective_pos:
                 entity_indices_to_keep.append(i)
         model_inputs["entity_ids"] = [model_inputs["entity_ids"][i] for i in entity_indices_to_keep]
-        model_inputs["entity_start_positions"] = [
-            model_inputs["entity_start_positions"][i] for i in entity_indices_to_keep
-        ]
-        model_inputs["entity_lengths"] = [model_inputs["entity_lengths"][i] for i in entity_indices_to_keep]
+        model_inputs["entity_position_ids"] = [model_inputs["entity_position_ids"][i] for i in entity_indices_to_keep]
 
         if self.entity_embeddings is not None:
-            model_inputs["entity_embeds"] = self.entity_embeddings[model_inputs["entity_ids"]]
+            model_inputs["entity_embeds"] = self.entity_embeddings[model_inputs["entity_ids"]].astype(np.float32)
         return model_inputs
 
     def __call__(self, text: str | list[str], **kwargs) -> BatchEncoding:
@@ -447,7 +449,9 @@ class KPRBertTokenizer(BertTokenizer):
         if isinstance(text, str):
             processed_inputs = self._preprocess_text(text, **kwargs)
             return BatchEncoding(
-                processed_inputs, tensor_type=kwargs.get("return_tensors", None), prepend_batch_axis=True
+                processed_inputs,
+                tensor_type=kwargs.get("return_tensors", None),
+                prepend_batch_axis=True,
             )
 
         processed_inputs_list: list[dict[str, list[int]]] = [self._preprocess_text(t, **kwargs) for t in text]
@@ -463,20 +467,33 @@ class KPRBertTokenizer(BertTokenizer):
                 return_attention_mask=kwargs.get("return_attention_mask"),
                 verbose=kwargs.get("verbose", True),
             )
-            # Collate entity_ids, entity_start_positions, and entity_lengths
+            # Pad entity ids
             max_num_entities = max(len(ids) for ids in collated_inputs["entity_ids"])
             for entity_ids in collated_inputs["entity_ids"]:
                 entity_ids += [0] * (max_num_entities - len(entity_ids))
-            for entity_start_positions in collated_inputs["entity_start_positions"]:
-                entity_start_positions += [-1] * (max_num_entities - len(entity_start_positions))
-            for entity_lengths in collated_inputs["entity_lengths"]:
-                entity_lengths += [0] * (max_num_entities - len(entity_lengths))
+            # Pad entity position ids
+            flattened_entity_length = [
+                len(ids) for ids_list in collated_inputs["entity_position_ids"] for ids in ids_list
+            ]
+            max_entity_token_length = max(flattened_entity_length) if flattened_entity_length else 0
+            for entity_position_ids_list in collated_inputs["entity_position_ids"]:
+                # pad entity_position_ids to max_entity_token_length
+                for entity_position_ids in entity_position_ids_list:
+                    entity_position_ids += [0] * (max_entity_token_length - len(entity_position_ids))
+                # pad to max_num_entities
+                entity_position_ids_list += [[0 for _ in range(max_entity_token_length)]] * (
+                    max_num_entities - len(entity_position_ids_list)
+                )
+            # Pad entity embeddings
             if "entity_embeds" in collated_inputs:
                 for i in range(len(collated_inputs["entity_embeds"])):
                     collated_inputs["entity_embeds"][i] = np.pad(
                         collated_inputs["entity_embeds"][i],
                         pad_width=(
-                            (0, max_num_entities - len(collated_inputs["entity_embeds"][i])),
+                            (
+                                0,
+                                max_num_entities - len(collated_inputs["entity_embeds"][i]),
+                            ),
                             (0, 0),
                         ),
                         mode="constant",