Prompt formatter API and canary transcribe tensor input support

[pzelasko]

What does this PR do ?

Generic prompt formatter for text modality with several out-of-the-box prompt format definitions. See the class documentation for more details.

Also, enables support for tensor/array inputs in Canary. Example snippet:

import os
from tempfile import NamedTemporaryFile

import lhotse

from nemo.collections.asr.models.aed_multitask_models import EncDecMultiTaskModel


def main():
    model = EncDecMultiTaskModel.from_pretrained("nvidia/canary-1b")

    path = ...
    rec = lhotse.Recording.from_file(path)
    audio = rec.load_audio()[0]

    TARGET_LANG = "es"

    # Array list input, legacy API for Canary-1B
    results = model.transcribe(
        audio=[audio],
        batch_size=1,
        num_workers=0,
        source_lang="en",
        target_lang=TARGET_LANG,
        task="asr",
        pnc="yes",
    )

    print(results)

    # Array list input, explicit single-turn prompt   
    results = model.transcribe(
        audio=[audio],
        batch_size=1,
        num_workers=0,
        role="user",
        slots={
            "source_lang": "en",
            "target_lang": TARGET_LANG,
            "task": "asr",
            "pnc": "yes",
        },
    )

    print(results)

    # Array list input, explicit multi-turn prompt   
    results = model.transcribe(
        audio=[audio],
        batch_size=1,
        num_workers=0,
        turns=[{
            "role": "user",
            "slots": {
                "source_lang": "en",
                "target_lang": TARGET_LANG,
                "task": "asr",
                "pnc": "yes",
            },
        }],
    )

    print(results)

    # Audio path input, explicit single-turn prompt   
    results = model.transcribe(
        audio=path,
        batch_size=1,
        num_workers=0,
        role="user",
        slots={
            "source_lang": "en",
            "target_lang": TARGET_LANG,
            "task": "asr",
            "pnc": "yes",
        },
    )

    print(results)

    # Legacy JSON manifest with slot values API for Canary-1B 
    with NamedTemporaryFile("w", suffix=".json") as f:
        lhotse.serialization.save_to_jsonl(
            [
                {
                    "audio_filepath": path,
                    "text": "irrelevant",
                    "duration": rec.duration,
                    "task": "asr",
                    "pnc": "yes",
                    "source_lang": "en",
                    "target_lang": TARGET_LANG,
                }
            ],
            f.name,
        )
        f.flush()
        os.fsync(f.fileno())

        results = model.transcribe(
            audio=f.name,
            batch_size=1,
            num_workers=0,
        )

        print(results)


if __name__ == "__main__":
    main()

We can also now provide these values dynamically to transcribe_speech.py from CLI. Example:

# Legacy Canary-1B format
python ~/code/NeMo/examples/asr/transcribe_speech.py \
  audio_dir=wavs \
  output_filename=out.json \
  batch_size=1 \
  pretrained_name=nvidia/canary-1b \
  +prompt.source_lang=en \
  +prompt.target_lang=es \
  +prompt.task=asr \
  +prompt.pnc=yes

# Explicit single-turn format
python ~/code/NeMo/examples/asr/transcribe_speech.py \
  audio_dir=wavs \
  output_filename=out.json \
  batch_size=1 \
  pretrained_name=nvidia/canary-1b \
  +prompt.role=user \
  +prompt.slots.source_lang=en \
  +prompt.slots.target_lang=es \
  +prompt.slots.task=asr \
  +prompt.slots.pnc=yes

# Explicit multi-turn format
python ~/code/NeMo/examples/asr/transcribe_speech.py \
  audio_dir=wavs \
  output_filename=out.json \
  batch_size=1 \
  pretrained_name=nvidia/canary-1b \
  +prompt.turns='[{role:user,slots:{source_lang:en,target_lang:es,task:asr,pnc:yes}}]'

Collection: ASR

Changelog

• Add specific line by line info of high level changes in this PR.

Usage

• You can potentially add a usage example below

# Add a code snippet demonstrating how to use this 

GitHub Actions CI

The Jenkins CI system has been replaced by GitHub Actions self-hosted runners.

The GitHub Actions CI will run automatically when the "Run CICD" label is added to the PR.
To re-run CI remove and add the label again.
To run CI on an untrusted fork, a NeMo user with write access must first click "Approve and run".

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  • Reviewer: Does the PR have correct import guards for all optional libraries?

PR Type:

• New Feature
• Bugfix
• Documentation

If you haven't finished some of the above items you can still open "Draft" PR.

Who can review?

Anyone in the community is free to review the PR once the checks have passed.
Contributor guidelines contains specific people who can review PRs to various areas.

Additional Information

• Related to # (issue)

[titu1994]

Initial comments

[stevehuang52]

Better change to raising TypeError and saying something like "expected input audio to have single channel", since users might not know what "MonoCut" means

[titu1994]

+1

[krishnacpuvvada]

If i understand correctly, this PR is for encoder-decoder models like canary/bestow, where all the (multi-turn) dialogue should be in text.
can we think little bit about supporting audio modality also in slot values. (may be we should keep audio slots untokenized and replace it with "audio features" later. one way is for prompt formatter to return something like list of lists)

[zhehuaichen]

+1 to this and Piotr told me he is planning on this as v2. maybe we can resume the discussion at the time

[krishnacpuvvada]

sg.
@pzelasko if possible, lets try to put the skeleton in place, e.g. if slot value needs to re-defined as (value, modality) tuple, return needs to be list of lists/tuples etc.

[pzelasko]

@krishnacpuvvada I'm thinking for multimodal we'll add a method that returns a "formatted prompt" as a sequence of embeddings instead. the benefit of using embeddings rather than token IDs is that we can support models with non-discrete latent spaces in addition to discretized. there are a few options:

1. initialize it with/register post-init a dict of {modality: nn.Module} that is used internally to convert "raw" modality input to sequence of embeddings; then, the prompt formatter is used at the beginning of forward step so that you can train these modules.
2. provide sequence of embeddings directly, but even then you still need to use the formatter in forward step as it's unlikely you'll embed audio/images/video in the dataloader process on a CPU fast enough.

in terms of skeletons, I've already put in the Modality type with a single type text that's used in slot schema definition and validation that a value "is" from a given modality. I 90% believe it'll be sufficient to extend to other modalities in V2.

[krishnacpuvvada]

sg.
Agreed. any audio encoder (especially our 600M ones) has to be run on GPU.

[zhehuaichen]

Nice work! LGTM

[krishnacpuvvada]

a thought.. can we also a sample.py/simple.py template with the simplest possible template and add few comments about which routines need to be defined. (This is mainly coming from - if a user wants to create their own custom template; I know there are plenty of examples already.. )

[titu1994]

I like this, yeah a canonical form of template to copy paste and directly modify

[pzelasko]

OK

[titu1994]

It looks really good now, minor comments from me, lets address the rest and merge

[titu1994]

+1

[titu1994]

Not important for this PR, but I was thinking of serializing the keys of the prompt format into config for user visibility.

[titu1994]

Very nice !

[titu1994]

I like this, yeah a canonical form of template to copy paste and directly modify

[titu1994]

Amazing work !

[stevehuang52]

Great work~! LGTM