refactor speechLM

Author: ryota-komatsu

README.md

@@ -62,11 +62,8 @@ audio_values = decoder(units)
 import torch
 import torchaudio
 from textless.data.speech_encoder import SpeechEncoder
-from tokenizers import Tokenizer
 from transformers import LlamaForCausalLM
 
-from src.speechlm.utils import convert_units_to_unicode
-
 wav_path = "/path/to/wav"
 
 encoder = SpeechEncoder.by_name(
@@ -77,21 +74,14 @@ encoder = SpeechEncoder.by_name(
     need_f0=False,
 ).cuda()
 
-# BPE tokenizer
-tokenizer = Tokenizer.from_file("/path/to/pretrained/tokenizer.json")
-
 model = LlamaForCausalLM.from_pretrained("/path/to/pretrained/model").cuda()
 
 # load a waveform
 waveform, sr = torchaudio.load(wav_path)
 waveform = torchaudio.functional.resample(waveform, sr, 16000)
 
 # encode a waveform into pseudo-phonetic units
-units = encoder(waveform.cuda())["units"].tolist()
-unicodes = convert_units_to_unicode(units)
-
-# BPE
-input_ids = tokenizer.encode(unicodes).ids
+input_ids = encoder(waveform.cuda())["units"].tolist()
 input_ids = torch.tensor([input_ids], device="cuda") + 2  # 0: pad, 1: EOS
 
 # Speech LM

configs/speechlm/hubert.yaml

@@ -2,14 +2,11 @@ dataset:
   wav_dir_train: "data/librilight"
   ext_audio: ".flac"
 
-  unicode_train: "data/speechlm/hubert/unicode/train"
-  train_file: "data/speechlm/hubert/unit/train.txt"
+  train: "ryota-komatsu/librilight"
   units_per_sample: 125
 
-  swuggy_dev_file: "data/speechlm/hubert/unit/lexical/dev.json"
-  sblimp_dev_file: "data/speechlm/hubert/unit/syntactic/dev.json"
-  swuggy_test_file: "data/speechlm/hubert/unit/lexical/test.json"
-  sblimp_test_file: "data/speechlm/hubert/unit/syntactic/test.json"
+  swuggy: "ryota-komatsu/swuggy"  # lexical
+  sblimp: "ryota-komatsu/sblimp"  # syntactic
 
   APP_DIR: "data/zr-data"
   result_dir: "results/speechlm/hubert"
@@ -19,7 +16,7 @@ dataloader:
 
 model:
   path: "models/speechlm/hubert"
-  vocab_size: 8192  # BPE vocab size
+  vocab_size: ${s2u.vocab_size}
   hidden_size: 768
   intermediate_size: 2048  # 4 * hidden_size * 2 / 3
   num_hidden_layers: 12
@@ -45,7 +42,4 @@ s2u:
   dense_model_name: "hubert-base-ls960"
   quantizer_model_name: "kmeans"
   vocab_size: 100
-
-  tokenizer_path: "models/speechlm/hubert/tokenizer.json"
-
   num_workers: 16

main_speechlm.py

@@ -2,18 +2,14 @@
 from omegaconf import OmegaConf
 
 from src.speechlm.eval import evaluate
-from src.speechlm.tokenize import encode, tokenize, tokenize_slm21
+from src.speechlm.tokenize import tokenize_slm21, tokenize_trainset
 from src.speechlm.train import train
 
 
 class TaskRunner:
-    def encode(self, config: str = "configs/speechlm/hubert.yaml", spkids: str = "1-9"):
+    def tokenize_trainset(self, config: str = "configs/speechlm/hubert.yaml"):
         config = OmegaConf.load(config)
-        encode(config, spkids)
-
-    def tokenize(self, config: str = "configs/speechlm/hubert.yaml"):
-        config = OmegaConf.load(config)
-        tokenize(config)
+        tokenize_trainset(config)
 
     def tokenize_slm21(self, config: str = "configs/speechlm/hubert.yaml"):
         config = OmegaConf.load(config)
@@ -29,8 +25,7 @@ def eval(self, config: str = "configs/speechlm/hubert.yaml"):
 
     def __call__(self, config: str = "configs/speechlm/hubert.yaml", spkids: str = "1-9"):
         config = OmegaConf.load(config)
-        encode(config, spkids)
-        tokenize(config)
+        tokenize_trainset(config, spkids)
         tokenize_slm21(config)
         train(config)
 

src/speechlm/data.py

@@ -1,5 +1,5 @@
 import random
-from typing import Any, Dict
+from typing import Any, Dict, Optional
 
 import torch
 import torchaudio
@@ -40,45 +40,37 @@ def collate_fn(batch):
         }
 
 
-class UnitDataset(torch.utils.data.Dataset):
-    def __init__(
-        self,
-        files,
-        units_per_sample: int,
-        num_special_tokens: int = 2,
-        eos_token_id: int = 1,
-    ):
-        self.input_ids = []
-        for file in files:
-            with open(file) as f:
-                for units in f:
-                    units = units.rstrip().split()
-                    units = torch.tensor([int(u) + num_special_tokens for u in units] + [eos_token_id])
-                    self.input_ids.append(units)
-
-        self.units_per_sample = units_per_sample
+def get_collate_fn(
+    num_special_tokens: int = 2,
+    pad_token_id: int = 0,
+    units_per_sample: Optional[int] = None,
+):
+    def collate_fn(batch) -> Dict[str, torch.LongTensor]:
+        input_ids = []
+        names = []
 
-    def __len__(self) -> int:
-        return len(self.input_ids)
+        for item in batch:
+            units = torch.tensor(item["units"]) + num_special_tokens
 
-    def __getitem__(self, n: int) -> Dict[str, torch.Tensor]:
-        input_ids = self.input_ids[n]
-        attention_mask = torch.ones_like(input_ids)
+            if units_per_sample:
+                diff = len(units) - units_per_sample
 
-        diff = len(input_ids) - self.units_per_sample
+                if diff > 0:
+                    start = random.randrange(diff)
+                    units = units[start : start + units_per_sample]
 
-        if diff > 0:
-            start = random.randrange(diff)
-            input_ids = input_ids[start : start + self.units_per_sample]
-            attention_mask = attention_mask[start : start + self.units_per_sample]
-        else:
-            input_ids = torch.nn.functional.pad(input_ids, (0, -diff))
-            attention_mask = torch.nn.functional.pad(attention_mask, (0, -diff))
+            input_ids.append(units)
+            names.append(item["id"])
 
-        labels = input_ids.masked_fill(input_ids.eq(0), -100)
+        input_ids = pad_sequence(input_ids, batch_first=True, padding_value=pad_token_id)
+        attention_mask = input_ids.ne(pad_token_id).long()
+        labels = input_ids.masked_fill(input_ids.eq(pad_token_id), -100)
 
         return {
             "input_ids": input_ids,
             "attention_mask": attention_mask,
             "labels": labels,
+            "names": names,
         }
+
+    return collate_fn

src/speechlm/eval.py

@@ -4,9 +4,10 @@
 import pandas as pd
 import torch
 import torch.nn.functional as F
+from datasets import load_dataset
 from transformers import LlamaForCausalLM
 
-from .utils import load_named_units_from_json
+from .data import get_collate_fn
 
 
 def evaluate(config):
@@ -22,17 +23,21 @@ def evaluate(config):
 
     _eval(
         model,
-        config.dataset.swuggy_test_file,
+        config.dataset.swuggy,
+        "test",
         Path(config.dataset.result_dir) / "lexical/test.txt",
         config.dataloader.batch_size_per_device,
         num_special_tokens,
+        config.model.pad_token_id,
     )
     _eval(
         model,
-        config.dataset.sblimp_test_file,
+        config.dataset.sblimp,
+        "test",
         Path(config.dataset.result_dir) / "syntactic/test.txt",
         config.dataloader.batch_size_per_device,
         num_special_tokens,
+        config.model.pad_token_id,
     )
 
     subprocess.run(
@@ -71,15 +76,24 @@ def evaluate(config):
 def _eval(
     model: LlamaForCausalLM,
     in_file,
+    split: str,
     out_file,
     batch_size: int,
     num_special_tokens: int = 2,
+    pad_token_id: int = 0,
 ):
+    dataset = load_dataset(in_file, split=split)
+    loader = torch.utils.data.DataLoader(
+        dataset,
+        batch_size,
+        collate_fn=get_collate_fn(num_special_tokens=num_special_tokens, pad_token_id=pad_token_id),
+    )
+
     with open(out_file, "w") as f:
-        for batch in load_named_units_from_json(in_file, batch_size, num_special_tokens):
+        for batch in loader:
             # Speech LM
             input_ids = batch["input_ids"].cuda()
-            labels = input_ids.masked_fill(input_ids.eq(0), -100)
+            labels = batch["labels"].cuda()
             logits = model(input_ids=input_ids, labels=labels).logits.transpose(1, 2)
 
             labels = F.pad(labels, (0, 1), value=-100)

src/speechlm/tokenize.py

@@ -1,38 +1,11 @@
-import glob
-import json
 from pathlib import Path
 
 import torch
+from datasets import Dataset, DatasetDict, Features, Sequence, Value
 from textless.data.speech_encoder import SpeechEncoder
-from tokenizers import Tokenizer
-from tokenizers.models import BPE
-from tokenizers.trainers import BpeTrainer
 from tqdm import tqdm
 
 from .data import SpeechDataset
-from .utils import convert_units_to_unicode, shift_unit
-
-
-def tokenize(config):
-    Path(config.s2u.tokenizer_path).parent.mkdir(parents=True, exist_ok=True)
-
-    files = glob.glob(config.dataset.unicode_train + "*")
-    initial_alphabet = [chr(shift_unit(unit)) for unit in range(config.s2u.vocab_size)]
-    trainer = BpeTrainer(vocab_size=config.model.vocab_size, initial_alphabet=initial_alphabet)
-    tokenizer = Tokenizer(BPE())
-    tokenizer.train(files=files, trainer=trainer)
-    tokenizer.save(config.s2u.tokenizer_path)
-
-    Path(config.dataset.train_file).parent.mkdir(parents=True, exist_ok=True)
-    with open(config.dataset.train_file, "w") as f:
-        for file in files:
-            with open(file) as g:
-                for unicodes in g:
-                    unicodes = unicodes.rstrip()
-                    units = tokenizer.encode(unicodes).ids
-                    units = " ".join(str(u) for u in units)
-
-                    f.write(f"{units}\n")
 
 
 def tokenize_slm21(config):
@@ -63,36 +36,41 @@ def tokenize_slm21(config):
         deduplicate=True,
         need_f0=False,
     ).cuda()
-    tokenizer = Tokenizer.from_file(config.s2u.tokenizer_path)
 
-    _tokenize_slm21(encoder, tokenizer, config.dataset.swuggy_dev_file, swuggy_dev_loader)
-    _tokenize_slm21(encoder, tokenizer, config.dataset.sblimp_dev_file, sblimp_dev_loader)
-    _tokenize_slm21(encoder, tokenizer, config.dataset.swuggy_test_file, swuggy_test_loader)
-    _tokenize_slm21(encoder, tokenizer, config.dataset.sblimp_test_file, sblimp_test_loader)
+    swuggy_dev = _tokenize(encoder, swuggy_dev_loader)
+    sblimp_dev = _tokenize(encoder, sblimp_dev_loader)
+    swuggy_test = _tokenize(encoder, swuggy_test_loader)
+    sblimp_test = _tokenize(encoder, sblimp_test_loader)
 
+    swuggy = DatasetDict({"dev": swuggy_dev, "test": swuggy_test})
+    sblimp = DatasetDict({"dev": sblimp_dev, "test": sblimp_test})
 
-def _tokenize_slm21(
+    swuggy.push_to_hub(config.dataset.swuggy)
+    sblimp.push_to_hub(config.dataset.sblimp)
+
+
+def _tokenize(
     encoder: SpeechEncoder,
-    tokenizer: Tokenizer,
-    file,
     data_loader: torch.utils.data.DataLoader,
 ):
-    Path(file).parent.mkdir(parents=True, exist_ok=True)
-
-    dataset = dict()
-
-    for item in tqdm(data_loader):
-        outputs = encoder(item["input_values"].cuda())
-        unicodes = convert_units_to_unicode(outputs["units"].tolist())
-        input_ids = tokenizer.encode(unicodes).ids
+    features = Features(
+        {
+            "id": Value("string"),
+            "units": Sequence(Value("int32")),
+        }
+    )
+
+    def generate_dataset():
+        for item in tqdm(data_loader):
+            outputs = encoder(item["input_values"].cuda())
+            units = outputs["units"].tolist()
 
-        dataset[item["name"][0]] = input_ids
+            yield {"id": item["name"][0], "units": units}
 
-    with open(file, "w") as f:
-        json.dump(dataset, f)
+    return Dataset.from_generator(generate_dataset, features=features)
 
 
-def encode(config, spk_ids: str = "1-9"):
+def tokenize_trainset(config, spk_ids: str = "1-9"):
     wav_dir_train = Path(config.dataset.wav_dir_train)
     train_paths = wav_dir_train.glob(f"*/[{spk_ids}]*/**/*" + config.dataset.ext_audio)
     train_set = SpeechDataset(train_paths)
@@ -106,15 +84,5 @@ def encode(config, spk_ids: str = "1-9"):
         need_f0=False,
     ).cuda()
 
-    _encode(encoder, config.dataset.unicode_train + f"{spk_ids}", train_loader)
-
-
-def _encode(encoder: SpeechEncoder, file, data_loader: torch.utils.data.DataLoader):
-    Path(file).parent.mkdir(parents=True, exist_ok=True)
-    with open(file, "w") as f:
-        for item in tqdm(data_loader):
-            outputs = encoder(item["input_values"].cuda())
-
-            unicodes = convert_units_to_unicode(outputs["units"].tolist())
-
-            f.write(f"{unicodes}\n")
+    trainset = _tokenize(encoder, train_loader)
+    trainset.push_to_hub(config.dataset.train, split=f"train{spk_ids}")