Rework handling of special tokens (#45)

docs/docusaurus_tsx/docs/FAQ/special_tokens.md

@@ -1,21 +1,39 @@
 # What special tokens are used?
 
-In the v2, special tokens were different for SEQ2SEQ and LM:
-LM was BOS, PAD, EOS with IDs (0, 1, 2) and the first vocab token started at id=3
-SEQ2SEQ was UNK, PAD, BOS, EOS with IDs (0, 1, 2, 3) and first vocab token started at id=4
-
-In v3 we changed this behavior to align things:
- group.add(
- "--default_specials",
- "-default_specilas",
- nargs="+",
- type=str,
- default=[
- DefaultTokens.UNK,
- DefaultTokens.PAD,
- DefaultTokens.BOS,
- DefaultTokens.EOS,
- ])
+There are 4 main special tokens:
+- BOS for "beginning of sentence";
+- PAD for "padding";
+- EOS for "end of sentence";
+- UNK for "unknown".
+
+## Special tokens actually used
+
+Depending on the context, these tokens can take various values:
+
+1. Default behaviour, training from scratch
+
+Some default values are defined as [constants](https://github.com/eole-nlp/eole/blob/ff39275c50d12951963008da11d029940b590713/eole/constants.py#L8) for the project:
+```python
+class DefaultTokens(object):
+ PAD = "<blank>"
+ BOS = "<s>"
+ EOS = "</s>"
+ UNK = "<unk>"
+```
+
+2. Retrieving a pretrained model from HF
+
+The special tokens will be retrieved and configured from the `special_tokens_map.json` configuration file from the HF model files.
+
+3. Custom behaviour
+
+In any case, these tokens can be overriden via the ad-hoc configuration settings:
+- `bos_token`
+- `pad_token`
+- `eos_token`
+- `unk_token`
+
+## Special tokens behaviour in Eole
 
 When we train a SEQ2SEQ model we use:
 SRC: srctok1 srctok2 srctok3 .... srctokn

eole/bin/convert/convert_HF.py

@@ -190,26 +190,6 @@
  "XLMRobertaXLForMaskedLM": TransformerEncoderModelConfig,
 }
 
-decoder_start_table = {
- "LlamaForCausalLM": "<s>",
- "MistralForCausalLM": "<s>",
- "MixtralForCausalLM": "<s>",
- "PhiForCausalLM": "",
- "Phi3ForCausalLM": "<s>",
- "GPT2LMHeadModel": "</s>",
- "XLMRobertaXLForMaskedLM": "<s>",
-}
-
-specials_table = {
- "LlamaForCausalLM": ["<unk>", "<s>", "</s>"],
- "MistralForCausalLM": ["<unk>", "<s>", "</s>"],
- "MixtralForCausalLM": ["<unk>", "<s>", "</s>"],
- "PhiForCausalLM": ["<unk>", "<s>", "</s>"],
- "Phi3ForCausalLM": ["<unk>", "<s>", "</s>"],
- "GPT2LMHeadModel": ["<unk>", "<s>", "</s>"],
- "XLMRobertaXLForMaskedLM": ["<s>", "<blank>", "</s>", "<unk>"],
-}
-
 
 class Tokenizer:
  def __init__(self, model_path: str):
@@ -313,6 +293,12 @@ def run(cls, args):
  )
  else:
  tokenizer_config_json = None
+ if os.path.exists(os.path.join(args.model_dir, "special_tokens_map.json")):
+ tokenizer_config_json = os.path.join(
+ args.model_dir, "special_tokens_map.json"
+ )
+ else:
+ tokenizer_config_json = None
  if os.path.exists(os.path.join(args.model_dir, "generation_config.json")):
  generation_config_json = os.path.join(
  args.model_dir, "generation_config.json"
@@ -415,6 +401,22 @@ def run(cls, args):
  raise huggingface_hub.utils.EntryNotFoundError(
  "No valid model files found"
  )
+ try:
+ try:
+ special_tokens_json = huggingface_hub.hf_hub_download(
+ repo_id=args.model_dir,
+ filename="special_tokens_map.json",
+ token=args.token,
+ )
+ except huggingface_hub.utils.EntryNotFoundError:
+ raise huggingface_hub.utils.EntryNotFoundError(
+ "Something went wrong the repo does not contain"
+ "any special_tokens_map.json file"
+ )
+ except Exception as e:
+ if isinstance(e, huggingface_hub.utils.EntryNotFoundError):
+ special_tokens_json = None
+ print(e)
 
  with open(config_path, encoding="utf-8") as fconfig:
  config = json.load(fconfig)
@@ -584,7 +586,6 @@ def run(cls, args):
  "n_positions": 0,
  }
  left_pad = True
- eos_token = None
  optional_eos = []
  mapped_tokens = []
  gpt2_pretok = False
@@ -947,18 +948,6 @@ def get_weight(checkpoint, tensor_name):
  for index in eos_token_id
  ]
  optional_eos = eos_tokens[1:]
- eos_token = eos_tokens[0]
- elif isinstance(eos_token_id, int):
- if "eos_token" in data.keys():
- if isinstance(data["eos_token"], dict):
- # Llama2 style
- eos_token = data["eos_token"]["content"]
- elif isinstance(data["eos_token"], str):
- eos_token = data["eos_token"]
- elif "added_tokens_decoder" in data.keys():
- eos_token = data["added_tokens_decoder"][str(eos_token_id)][
- "content"
- ]
  # Automatically convert added_tokens into mapped_tokens
  if "added_tokens_decoder" in data.keys():
  mapped_tokens = [
@@ -973,6 +962,29 @@ def get_weight(checkpoint, tensor_name):
  else:
  add_bos_token = True
 
+ vocabs = {"specials": {}}
+
+ if special_tokens_json is not None:
+ with open(special_tokens_json, encoding="utf-8") as f:
+ special_tokens_map = json.load(f)
+ for token_name in ["bos_token", "unk_token", "eos_token", "pad_token"]:
+ token = special_tokens_map.get(token_name, None)
+ if isinstance(token, list):
+ vocabs["specials"][token_name] = token[0]
+ elif isinstance(token, str):
+ vocabs["specials"][token_name] = token
+ elif isinstance(token, dict):
+ vocabs["specials"][token_name] = token["content"]
+ elif tokenizer_json is not None:
+ with open(tokenizer_json, encoding="utf-8") as f:
+ data = json.load(f)
+ vocab = {v: k for k, v in data["model"]["vocab"].items()}
+ for token_name in ["bos_token", "unk_token", "eos_token", "pad_token"]:
+ if f"{token_name}_id" in config.keys():
+ vocabs["specials"][token_name] = vocab[
+ config[f"{token_name}_id"]
+ ]
+
  if generation_config_json is not None:
  with open(generation_config_json, encoding="utf-8") as f:
  data = json.load(f)
@@ -982,8 +994,6 @@ def get_weight(checkpoint, tensor_name):
  for key in keys:
  if key in data.keys():
  generation_config_dict[key] = data[key]
-
- vocabs = {}
  if (
  tokenizer_model is not None
  ): # sentencepiece mode (might be good to check it's a SP model)
@@ -1003,26 +1013,19 @@ def get_weight(checkpoint, tensor_name):
  vocab.extend(newtokens)
  for tok in data["added_tokens"]:
  vocab[tok["id"]] = tok["content"]
- if "<|startoftext|>" in vocab:
- index = vocab.index("<|startoftext|>")
- vocab[index] = DefaultTokens.BOS
- if eos_token is not None:
- if eos_token in vocab and "</s>" not in vocab:
- index = vocab.index(eos_token)
- vocab[index] = DefaultTokens.EOS
- if "<0x00>" in vocab:
- index = vocab.index("<0x00>")
- vocab[index] = DefaultTokens.PAD
  src_vocab = pyonmttok.build_vocab_from_tokens(
  vocab,
- special_tokens=specials_table[arch],
  )
  else: # # BPE mode - we leverage the HF tokenizer.json info
  src_subword_type = "bpe"
  with open(tokenizer_json, encoding="utf-8") as f:
  data = json.load(f)
  # gpt2_pretok
  pretokenizers = data.get("pre_tokenizer", {}).get("pretokenizers", [{}])
+ pre_tokenizer = data.get("pre_tokenizer", None)
+ pretokenizers = pre_tokenizer.get("pretokenizers", None)
+ if pretokenizers is None:
+ pretokenizers = [pre_tokenizer]
  for pretokenizer in pretokenizers:
  if pretokenizer.get("type", None) == "ByteLevel":
  gpt2_pretok = True
@@ -1031,23 +1034,14 @@ def get_weight(checkpoint, tensor_name):
  # "Ā" is '\x00' in unicode (cf tokenize.py gpt2 mapping)
  for tok in data["model"]["vocab"]
  ]
+ if DefaultTokens.PAD in vocab:
+ vocabs["specials"]["pad_token"] = DefaultTokens.PAD
  voc_size = len(vocab)
  if vocab_size > voc_size:
  for i in range(vocab_size - voc_size):
  vocab.append(DefaultTokens.VOCAB_PAD + str(i))
  for tok in data["added_tokens"]:
  vocab[tok["id"]] = tok["content"]
- if "<|startoftext|>" in vocab:
- index = vocab.index("<|startoftext|>")
- vocab[index] = DefaultTokens.BOS
- if "<|begin_of_text|>" in vocab:
- index = vocab.index("<|begin_of_text|>")
- vocab[index] = DefaultTokens.BOS
- if eos_token is not None:
- if eos_token in vocab and "</s>" not in vocab:
- index = vocab.index(eos_token)
- vocab[index] = DefaultTokens.EOS
-
  src_vocab = pyonmttok.build_vocab_from_tokens(vocab)
 
  tokenizer_basename = "bpe.model"
@@ -1062,7 +1056,7 @@ def get_weight(checkpoint, tensor_name):
  vocabs["src"] = src_vocab
  vocabs["tgt"] = src_vocab
  if add_bos_token:
- vocabs["decoder_start_token"] = decoder_start_table[arch]
+ vocabs["decoder_start_token"] = vocabs["specials"]["bos_token"]
  else:
  vocabs["decoder_start_token"] = ""
  vocab_dict = vocabs_to_dict(vocabs)
@@ -1089,6 +1083,7 @@ def get_weight(checkpoint, tensor_name):
  tgt_vocab_size=vocab_size,
  vocab_size_multiple=8,
  decoder_start_token=vocabs["decoder_start_token"],
+ **vocabs["specials"],
  transforms=["onmt_tokenize", "filtertoolong"],
  transforms_configs={
  "filtertoolong": {"src_seq_length": 512, "tgt_seq_length": 512},

eole/bin/tools/LM_scoring.py

@@ -82,7 +82,8 @@ def run(cls, args):
  )
 
  vocabs, model, model_opt = config.model.model_class.load_test_model(config)
- padding_idx = vocabs["tgt"][DefaultTokens.PAD]
+ pad_token = vocabs["specials"].get("pad_token", DefaultTokens.PAD)
+ padding_idx = vocabs["tgt"][pad_token]
  criterion = torch.nn.CrossEntropyLoss(
  ignore_index=padding_idx, reduction="none"
  )

eole/config/data.py

@@ -26,15 +26,17 @@ class BaseVocabConfig(Config):
  description="Default decoder start token. For most models it is <s> = BOS. "
  "Some fairseq models require </s>.",
  )
- default_specials: list = Field(
- default=[
- constants.DefaultTokens.UNK,
- constants.DefaultTokens.PAD,
- constants.DefaultTokens.BOS,
- constants.DefaultTokens.EOS,
- ],
- description="Default specials used for vocab initialization. "
- "UNK, PAD, BOS, EOS will take IDs 0, 1, 2, 3.",
+ bos_token: str | None = Field(
+ default=constants.DefaultTokens.BOS,
+ )
+ eos_token: str | None = Field(
+ default=constants.DefaultTokens.EOS,
+ )
+ unk_token: str | None = Field(
+ default=constants.DefaultTokens.UNK,
+ )
+ pad_token: str | None = Field(
+ default=constants.DefaultTokens.PAD,
  )
  # pre trained embeddings stuff, might be put elsewhere
  both_embeddings: str | None = Field(

eole/inference_engine.py

@@ -278,6 +278,13 @@ def __init__(self, config, model_type=None):
  vocabs["src"] = src_vocab
  vocabs["tgt"] = src_vocab
  vocabs["decoder_start_token"] = "<s>"
+ # TODO: this should be loaded from model config
+ vocabs["specials"] = {
+ "bos_token": DefaultTokens.BOS,
+ "pad_token": DefaultTokens.PAD,
+ "eos_token": DefaultTokens.EOS,
+ "unk_token": DefaultTokens.UNK,
+ }
  self.vocabs = vocabs
  # Build transform pipe
  transforms = make_transforms(config, self.transforms_cls, self.vocabs)
@@ -290,7 +297,10 @@ def predict_batch(self, batch, config):
  _input_tokens = [
  self.vocabs["src"].lookup_index(id)
  for id in start_ids
- if id != self.vocabs["src"].lookup_token(DefaultTokens.PAD)
+ if id
+ != self.vocabs["src"].lookup_token(
+ self.vocabs["specials"].get("pad_token", DefaultTokens.PAD)
+ )
  ]
  input_tokens.append(_input_tokens)
  if self.model_type == ModelType.DECODER:

eole/inputters/inputter.py

@@ -16,6 +16,13 @@ def build_vocab(config, specials):
  'decoder_start_token': DefaultTokens.BOS
  }
  """
+ vocabs = {}
+ vocabs["specials"] = {
+ "bos_token": config.bos_token,
+ "pad_token": config.pad_token,
+ "eos_token": config.eos_token,
+ "unk_token": config.unk_token,
+ }
 
  def _pad_vocab_to_multiple(vocab, multiple):
  vocab_size = len(vocab)
@@ -26,8 +33,8 @@ def _pad_vocab_to_multiple(vocab, multiple):
  vocab.add_token(DefaultTokens.VOCAB_PAD + str(i))
  return vocab
 
- default_specials = config.default_specials
- vocabs = {}
+ default_specials = list(vocabs["specials"].values())
+
  src_vocab = _read_vocab_file(config.src_vocab, config.src_words_min_frequency)
 
  src_specials = [
@@ -45,7 +52,7 @@ def _pad_vocab_to_multiple(vocab, multiple):
  src_vocab = pyonmttok.build_vocab_from_tokens(
  src_vocab, maximum_size=config.src_vocab_size, special_tokens=src_specials
  )
- src_vocab.default_id = src_vocab[DefaultTokens.UNK]
+ src_vocab.default_id = src_vocab[config.unk_token]
  if src_vocab.default_id >= len(src_vocab):
  src_vocab.default_id = (
  0 # patch that assigns OOV to id=0 when UNK does not exist
@@ -80,7 +87,6 @@ def _pad_vocab_to_multiple(vocab, multiple):
  vocabs["tgt"] = tgt_vocab
 
  vocabs["decoder_start_token"] = config.decoder_start_token
-
  return vocabs
 
 
@@ -126,6 +132,8 @@ def vocabs_to_dict(vocabs):
  vocabs_dict["decoder_start_token"] = vocabs["decoder_start_token"]
  else:
  vocabs_dict["decoder_start_token"] = DefaultTokens.BOS
+ if "specials" in vocabs.keys():
+ vocabs_dict["specials"] = vocabs["specials"]
  return vocabs_dict
 
 
@@ -148,5 +156,7 @@ def dict_to_vocabs(vocabs_dict):
  vocabs["tgt"] = pyonmttok.build_vocab_from_tokens(vocabs_dict["tgt"])
  if vocabs["tgt"].default_id >= len(vocabs["src"]):
  vocabs["tgt"].default_id = 0 # patch that assigns OOV to id=0
+ if "specials" in vocabs_dict.keys():
+ vocabs["specials"] = vocabs_dict["specials"]
 
  return vocabs