On punctuation and capitalization

[1-800-BAD-CODE]

Here's a few feature requests + bugs related to the punctuation and capitalization model.

Punctuation issues

Inverted punctuation

For languages like Spanish, we need two predictions per token to account for the possibility of inverted punctuation tokens preceding a word.

Subword masking

By always applying subtoken masks, continuous-script languages (e.g., Chinese) cannot be punctuated without applying some sort of pre-processing. It would be useful if the model processed text in its native script.

Arbitrary punctuation tokens

The text-based data set does not allow to punctuate languages such as Thai, where a space character is a punctuation token. These languages could work by having a token-based data set and removing subword masks (essentially, resolving the other issues resolves this one).

Capitilization issues

Acronyms

The capitalization prediction is simply whether a word starts with a capital letter. So acronyms like 'amc' will not be correctly capitalized.

Names that begin with a particle

Similar to the acronym issue, words that begin with a particle, e.g., 'mcdonald', cannot be properly capitalized to 'McDonald'.

Capitalization is independent of punctuation

Currently, the two heads are conditioned only on the encoder's output and independent of each other. But capitalization is dependent on punctuation in many cases.

An example of what might go wrong is we may end up with "Hello, world, What's up?" because the capitalization model might expect a period after 'world'. Essentially the capitalization head is predicting what the punctuation head will do.

In practice I have found this problem to be an uncommon manifestation, but to be correct, capitalization should take into account the punctuator's output. Implicitly, we are forcing the capitalization head to learn punctuation (and predict the punctuation head's output).

Data set issues

Dataset as text

First (this may be a personal preference), it is unnatural to have a preprocessed data set in text format rather than token IDs.
More importantly, text data sets are incompatible with other issues mentioned in this thread (subwords, space as punctuation)

Supported data set classes

Dataset is fixed to one class, but it would be more convenient to simply expect an abstract base class, and let the user specify a _target_ in the config and use hydra.utils.instantiate, as in some other models. E.g.,

NeMo/nemo/collections/tts/models/melgan.py

Line 298 in bc6215f

dataset = instantiate(cfg.dataset)

For example, a user may wish to implement a different dataset that generates examples on-the-fly, or use a ConcatDataset with multiple languages and temperature sampling, etc.

Paragraph segmentation

A primary benefit of this model is to improve NMT results on unpunctuated and uncased ASR output. However, running MT on an arbitrarily-long inputs will inevitably end poorly.

For this model to be complete, I would argue it needs to implement a third token classification analytic: paragraph segmentation (splitting a paragraph into its constituent sentences). Translating sentences as separate units would improve results in many cases. Furthermore, a Transformer's runtime complexity is N^2 in the sequence length.

[1-800-BAD-CODE]

I'm willing to implement and contribute the following solution if there is sufficient interest.

This solution should generalize to any language and solve all problems mentioned above. Given that it is vastly different from the current punctuation + capitalization model, I would propose a separate model (e.g. TextStructuringModel) rather than trying to change the existing model.

Here's a little bit of reasoning on which this solution is founded:

---

• Premise 1. Subword tokenizers should be preferred over character tokenizers.
• Premise 2. Capitalization requires a character tokenizer.
• Implication 1: Two tokenizers should be used: subword for punctuation/segmentation, character for capitalization.
• Implication 2: The model will use two language models/encoders, one for each tokenizer. [Maybe the vocab can be combined and use a single LM.]

---

• Premise 1. The model should generalize well to all languages and have minimal language-specific paths.
• Premise 2: It's OK to make unnecessary predictions if they are reasonably cheap.
• Implication 1: The punctuation head should always make two predictions per token, to account for Spanish, Asturian, etc. The model will learn that other languages have no punctuation before tokens and simply predict null.
• Implication 2: The model should never use a subword mask, to allow continuous-script languages to be punctuated and segmented between subword units. The model will learn that non-continuous languages are never broken or punctuated between subword units.

---

• Premise 1: Each analytic should be a function of the output of all dependent analytics.
• Implication 1: At inference time, the analytics are chained together such that punctuation sees unstructured text, capitalization sees punctuated text, and segmentation sees capitalized and punctuated text.
• Implication 2: At training time, the data loader produces 3 different versions of the inputs to provide reference data w.r.t. the particular analytic. Essentially, we are using teacher forcing to provide the reference output of dependent analytics.

---

The model's components are shown here, along with the training algorithm:

---

At inference time, it would run like this:

[ekmb]

Hi @1-800-BAD-CODE, thank you for bringing up these valid issues and your willingness to contribute. It would be great if you could add a model that doesn't rely on subword masking with the punctuation head that predicts pre-and post- punctuation marks. I agree that having a separate model would make more sense, and we can then re-factor/depreciate the current model.

A few questions/comments:

1. If a training example for the punctuation model represents 1+ sentences, the model learns to split them into separate sentences, i.e., put the end of sentence punctuation marks in between the sentences. Paragraph segmentation seems redundant. Please also see this method that removes 'margin' probabilities when combining punctuation predictions from multiple segments (This was specifically added to segment long input texts for NMT task.)
2. I agree that capitalization should depend on punctuation. As for char-based tokenization, using a single LM would be preferable. The vocabularies of the most pre-trained transformer models include single chars (see BERT vocab.txt). However, the chars don't carry any semantic meaning, and the input sequence is getting longer (this will affect the segmentation head if you still want it). An alternative solution could be to use the same tokenization for the subword task, add punctuation head output, similarly to [this] (https://assets.amazon.science/55/3f/8c27b4014bdd983087fdb1d73412/robust-prediction-of-punctuation-and-truecasing-for-medical-asr.pdf) and use capitalization tags: {first_letter_upper, all_caps, all_lower, mixed (for McDonald cases)}.

[1-800-BAD-CODE]

Thanks for the feedback. I will implement something that uses only one LM and solves all problems in my original post, so not exactly like my initial plan outlined above.

A couple notes (anyone's feedback would be useful, but not necessary):

1. Based on my understanding of that paper (seemingly the most-cited deep-learning approach to this problem), they simply predicted whether a word was "mixed case" or not, without predicting which characters should be upper/lower case.
   My initial thought to solve this problem, utilizing a subword model, is to have the capitalization head makes N binary decisions per token, where N is the maximum subword length (default 16 for sentencepiece). Prediction i is whether subword letter i is capitalized. Whether this is feasible might best be answered experimentally.

2. Based on my understanding of that paper, capitalization for token i is conditioned only the predicted punctuation of token i. In practice, capitalization is usually dependent on token i ("Dr.") or i-1 (new sentence). But I would say that the only sound conditioning method is to allow the LM to encode the newly-punctuated sequence, although at inference that requires multiple, sequential passes of the LM per query.

[ekmb]

Hi @1-800-BAD-CODE, any updates?

[1-800-BAD-CODE]

It's getting there. It's been unexpectedly-difficult to tune the punctuation loss weights, which is exacerbated by allowing predictions between subword tokens and predicting twice per token, greatly increasing the proportion of null class predictions.

I hope to have a more meaningful update with a couple of weeks.

[itzsimpl]

Hi @1-800-BAD-CODE, any updates? Very interested in this feature.

[1-800-BAD-CODE]

It's still a work in progress, but active.

[1-800-BAD-CODE]

It lives. This model was fine-tuned from bert-base-multilingual-cased for ~30 minutes on Spanish, English, Russian, and Chinese.

Still a lot of things to smooth out.

>>> from nemo.collections.nlp.models.token_classification.punct_cap_seg_model import PunctCapSegModel

>>> m = PunctCapSegModel.restore_from("pcs_esruenzh_bert_base_multilingual_cased.nemo", map_location="cpu")

>>> sentences = [
    "milk was a bad choice it's so hot outside", 
    "молоко было плохим выбором на улице так жарко",    # same sentence in Russian
    "la leche fue una mala elección hace mucho calor afuera",   # same sentence in Spanish
    "牛奶是个糟糕的选择外面非常热",   # same sentence in Chinese
    "no he perdonado ni a ti ni a maria por qué no juegas conmigo"  # Spanish sentence with question in the second half
]

>>> results: List[List[str]] = m.infer(sentences)

>>> for input_sentence, outputs in zip(sentences, results):
...   print(f"Input: {input_sentence}")
...   for output in outputs:
...     print(f"  Output: {output}")
... 
Input: milk was a bad choice it's so hot outside
  Output: Milk was a bad choice.
  Output: It's so hot outside.
Input: молоко было плохим выбором на улице так жарко
  Output: Молоко было плохим выбором.
  Output: На улице так жарко. 
Input: la leche fue una mala elección hace mucho calor afuera
  Output: La leche fue una mala elección.
  Output: Hace mucho calor afuera.
Input: 牛奶是个糟糕的选择外面非常热
  Output: 牛 奶 是 个 糟 糕 的 选 择 。
  Output: 外 面 非 常 热 。
Input: no he perdonado ni a ti ni a maria por qué no juegas conmigo
  Output: No he perdonado ni a ti ni a Maria.
  Output: ¿ Por qué no juegas conmigo?

[1-800-BAD-CODE]

See pull #4637