Update LLaST implementation

.pre-commit-config.yaml

@@ -4,7 +4,7 @@ repos:
     rev: 5.0.4
     hooks:
       - id: flake8
-        args: ["--exclude=xtuner/model/transformers_models/*"]
+        args: ["--exclude=xtuner/model/transformers_models/*,xtuner/evaluation/metrics/sacrebleu.py"]
   - repo: https://github.com/PyCQA/isort
     rev: 5.11.5
     hooks:

xtuner/dataset/__init__.py

@@ -7,6 +7,7 @@
                           load_intern_repo_tokenized_dataset,
                           load_intern_repo_untokenized_dataset)
 from .json_dataset import load_json_file
+from .llast import LLaSTDataset
 from .llava import LLaVADataset
 from .modelscope import process_ms_dataset
 from .moss_sft import MOSSSFTDataset
@@ -24,5 +25,5 @@
     'load_intern_repo_tokenized_dataset',
     'load_intern_repo_untokenized_dataset', 'build_packed_dataset',
     'RefCOCOJsonDataset', 'RefCOCOJsonEvalDataset', 'InvRefCOCOJsonDataset',
-    'load_json_file'
+    'load_json_file', 'LLaSTDataset'
 ]

xtuner/dataset/collate_fns/__init__.py

@@ -1,5 +1,8 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 from .default_collate_fn import default_collate_fn
+from .llast_collate_fn import llast_audiomask_mel_collate_fn
 from .mmlu_collate_fn import mmlu_collate_fn
 
-__all__ = ['default_collate_fn', 'mmlu_collate_fn']
+__all__ = [
+    'default_collate_fn', 'mmlu_collate_fn', 'llast_audiomask_mel_collate_fn'
+]

xtuner/dataset/collate_fns/llast_collate_fn.py

@@ -0,0 +1,60 @@
+# Copyright (c) LLaST. All rights reserved.
+# Copyright (c) OpenMMLab. All rights reserved.
+from typing import Dict, Sequence
+
+import torch
+from torch.nn.utils.rnn import pad_sequence
+
+from xtuner.utils import (DEFAULT_PAD_TOKEN_INDEX, IGNORE_INDEX,
+                          LLAST_AUDIO_PADDING_TOKEN_INDEX)
+
+
+def llast_audiomask_mel_collate_fn(
+        instances: Sequence[Dict],
+        pad_index: int = DEFAULT_PAD_TOKEN_INDEX,
+        return_hf_format: bool = False) -> Dict[str, torch.Tensor]:
+    """Add audio tokens and conduct padding operation."""
+    input_ids = []
+    labels = []
+    feats_lens = []
+    has_audio = any(inst.get('audio_tokens') is not None for inst in instances)
+
+    if has_audio:
+        audio_tokens = []
+    for example in instances:
+        input_ids.append(torch.tensor(example['input_ids']))
+        labels.append(torch.tensor(example['labels']))
+        if has_audio:
+            audio_tokens.append(example['audio_tokens'])
+        feats_lens.append(torch.tensor(example['audio_lens']))
+    if len(instances) > 1:
+        input_ids = pad_sequence(
+            input_ids, batch_first=True, padding_value=pad_index)
+        labels = pad_sequence(
+            labels, batch_first=True, padding_value=IGNORE_INDEX)
+        # padding audio tokens
+        padded_audio_tokens = pad_sequence(
+            audio_tokens,
+            batch_first=True,
+            padding_value=LLAST_AUDIO_PADDING_TOKEN_INDEX)
+
+    else:
+        input_ids = torch.stack(input_ids)
+        labels = torch.stack(labels)
+        padded_audio_tokens = torch.stack(audio_tokens)
+
+    data_dict = {
+        'input_ids': input_ids,
+        'attention_mask': input_ids.ne(pad_index),
+        'labels': labels
+    }
+
+    if has_audio:
+        audio_lens = torch.stack(feats_lens)
+        data_dict['audio_tokens'] = padded_audio_tokens
+        data_dict['audio_lens'] = audio_lens
+
+    if return_hf_format:
+        return data_dict
+    else:
+        return {'data': data_dict, 'data_samples': instances}

xtuner/dataset/huggingface.py

@@ -65,8 +65,8 @@ def add_template_to_dataset(dataset, template_map_fn, map_num_proc):
 
 
 def tokenize_dataset(dataset, tokenizer, max_length, with_image_token,
-                     input_ids_with_output, remove_unused_columns,
-                     map_num_proc):
+                     with_audio_token, input_ids_with_output,
+                     remove_unused_columns, map_num_proc):
     assert (tokenizer is not None) and (max_length is not None), \
         f'({tokenizer}, {max_length})'
     if isinstance(tokenizer, dict) or isinstance(
@@ -78,6 +78,7 @@ def tokenize_dataset(dataset, tokenizer, max_length, with_image_token,
             tokenizer=tokenizer,
             max_length=max_length,
             with_image_token=with_image_token,
+            with_audio_token=with_audio_token,
             input_ids_with_output=input_ids_with_output),
         remove_columns=list(dataset.column_names)
         if remove_unused_columns else None,
@@ -112,6 +113,7 @@ def process(dataset,
             use_varlen_attn=False,
             input_ids_with_output=True,
             with_image_token=False,
+            with_audio_token=False,
             map_num_proc=32):
     """Post-process the dataset loaded from the Hugging Face Hub, or a local
     dataset.
@@ -153,6 +155,9 @@ def process(dataset,
         with_image_token: Whether to convert DEFAULT_IMAGE_TOKEN to
             IMAGE_TOKEN_INDEX. Typically set it to True during the training
             of VLM.
+        with_audio_token: Whether to convert DEFAULT_AUDIO_TOKEN to
+            LLAST_AUDIO_TOKEN_INDEX. Typically set it to True during the
+            training of SLM.
         map_num_proc: Max number of processes when mapping the dataset.
     """
     if use_varlen_attn:
@@ -197,7 +202,8 @@ def process(dataset,
 
     if do_dataset_tokenization:
         dataset = tokenize_dataset(dataset, tokenizer, max_length,
-                                   with_image_token, input_ids_with_output,
+                                   with_image_token, with_audio_token,
+                                   input_ids_with_output,
                                    remove_unused_columns, map_num_proc)
 
     if input_ids_with_output:
@@ -213,7 +219,7 @@ def process(dataset,
                                shuffle_before_pack, map_num_proc)
 
     # add 'length'
-    dataset = dataset.map(get_lengths, num_proc=map_num_proc)
+    dataset = dataset.map(get_lengths, num_proc=1)
     setattr(dataset, 'length', dataset['length'])
 
     return dataset
@@ -234,6 +240,7 @@ def process_hf_dataset(dataset,
                        use_varlen_attn=False,
                        input_ids_with_output=True,
                        with_image_token=False,
+                       with_audio_token=False,
                        map_num_proc=32):
     """Post-process the dataset loaded from the Hugging Face Hub, or a local
     dataset.
@@ -275,6 +282,9 @@ def process_hf_dataset(dataset,
         with_image_token: Whether to convert DEFAULT_IMAGE_TOKEN to
             IMAGE_TOKEN_INDEX. Typically set it to True during the training
             of VLM.
+        with_audio_token: Whether to convert DEFAULT_AUDIO_TOKEN to
+            LLAST_AUDIO_TOKEN_INDEX. Typically set it to True during the
+            training of SLM.
         map_num_proc: Max number of processes when mapping the dataset.
     """
     kwargs = dict(
@@ -293,6 +303,7 @@ def process_hf_dataset(dataset,
         use_varlen_attn=use_varlen_attn,
         input_ids_with_output=input_ids_with_output,
         with_image_token=with_image_token,
+        with_audio_token=with_audio_token,
         map_num_proc=map_num_proc)
     if not (dist.is_available() and dist.is_initialized()):
         return process(**kwargs)