init
This commit is contained in:
smthemex
@@ -0,0 +1,84 @@
|
||||
# Copyright (c) ByteDance, Inc. and its affiliates.
|
||||
# Copyright (c) Chutong Meng
|
||||
#
|
||||
# This source code is licensed under the CC BY-NC license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
# Based on AudioDec (https://github.com/facebookresearch/AudioDec)
|
||||
|
||||
import torch.nn as nn
|
||||
|
||||
from .modules.decoder import Decoder
|
||||
from .modules.encoder import Encoder
|
||||
from .modules.projector import Projector
|
||||
from .modules.quantizer import Quantizer
|
||||
|
||||
|
||||
class RepCodec(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
input_channels=768,
|
||||
output_channels=768,
|
||||
encode_channels=768,
|
||||
decode_channels=768,
|
||||
code_dim=768,
|
||||
codebook_num=1,
|
||||
codebook_size=1024,
|
||||
bias=True,
|
||||
enc_ratios=(1, 1),
|
||||
dec_ratios=(1, 1),
|
||||
enc_strides=(1, 1),
|
||||
dec_strides=(1, 1),
|
||||
enc_kernel_size=3,
|
||||
dec_kernel_size=3,
|
||||
enc_block_dilations=(1, 1),
|
||||
enc_block_kernel_size=3,
|
||||
dec_block_dilations=(1, 1),
|
||||
dec_block_kernel_size=3
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.input_channels = input_channels
|
||||
|
||||
self.encoder = Encoder(
|
||||
input_channels=input_channels,
|
||||
encode_channels=encode_channels,
|
||||
channel_ratios=enc_ratios,
|
||||
strides=enc_strides,
|
||||
kernel_size=enc_kernel_size,
|
||||
bias=bias,
|
||||
block_dilations=enc_block_dilations,
|
||||
unit_kernel_size=enc_block_kernel_size
|
||||
)
|
||||
|
||||
self.decoder = Decoder(
|
||||
code_dim=code_dim,
|
||||
output_channels=output_channels,
|
||||
decode_channels=decode_channels,
|
||||
channel_ratios=dec_ratios,
|
||||
strides=dec_strides,
|
||||
kernel_size=dec_kernel_size,
|
||||
bias=bias,
|
||||
block_dilations=dec_block_dilations,
|
||||
unit_kernel_size=dec_block_kernel_size
|
||||
)
|
||||
|
||||
self.projector = Projector(
|
||||
input_channels=self.encoder.out_channels,
|
||||
code_dim=code_dim,
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
bias=False
|
||||
)
|
||||
|
||||
self.quantizer = Quantizer(
|
||||
code_dim=code_dim,
|
||||
codebook_num=codebook_num,
|
||||
codebook_size=codebook_size
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.encoder(x)
|
||||
z = self.projector(x)
|
||||
zq, vqloss, perplexity = self.quantizer(z)
|
||||
y = self.decoder(zq)
|
||||
return y, zq, z, vqloss, perplexity
|
||||
@@ -0,0 +1,212 @@
|
||||
# Copyright (c) ByteDance, Inc. and its affiliates.
|
||||
# Copyright (c) Chutong Meng
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import argparse
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Tuple, List, Optional
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import yaml
|
||||
from tqdm import tqdm
|
||||
|
||||
from .RepCodec import RepCodec
|
||||
|
||||
ALL_MODELS = {
|
||||
"data2vec_base_l6": 768,
|
||||
"data2vec_large_l18": 1024,
|
||||
"hubert_base_l9": 768,
|
||||
"hubert_large_l18": 1024,
|
||||
"whisper_medium_l24": 1024,
|
||||
"whisper_large_l32": 1280
|
||||
}
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
||||
parser.add_argument(
|
||||
"in_dir",
|
||||
type=str,
|
||||
help="directory of representations to be tokenized."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
required=True,
|
||||
type=str,
|
||||
help="path of the RepCodec model."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tsv_path",
|
||||
required=True,
|
||||
type=str,
|
||||
help="path of the tsv file."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model_config_path",
|
||||
default=None,
|
||||
type=str,
|
||||
help="please provide this training config if you are using the model you trained yourself."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--n_shard",
|
||||
required=False,
|
||||
type=int,
|
||||
default=1,
|
||||
help="number of shards of representations."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--use_gpu",
|
||||
default=False,
|
||||
action="store_true",
|
||||
help="whether use gpu for inference."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--batch_size",
|
||||
default=1,
|
||||
type=int,
|
||||
help="number of utterances for each mini batch."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--out_dir",
|
||||
type=str,
|
||||
default=".",
|
||||
help="the directory to save the output."
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def load_model(model_path: str, config_path: Optional[str] = None):
|
||||
if config_path is None:
|
||||
name = os.path.basename(model_path).strip(".pkl")
|
||||
assert name in ALL_MODELS.keys(), f"Cannot find configs for {model_path}. " \
|
||||
f"Please provide the config file you used for training."
|
||||
config = os.path.join(os.path.dirname(__file__), "configs", f"repcodec_dim{ALL_MODELS[name]}.yaml")
|
||||
with open(config) as fp:
|
||||
conf = yaml.load(fp, Loader=yaml.FullLoader)
|
||||
else:
|
||||
with open(config_path) as fp:
|
||||
conf = yaml.load(fp, Loader=yaml.FullLoader)["model_params"]
|
||||
|
||||
model = RepCodec(**conf)
|
||||
model.load_state_dict(torch.load(model_path, map_location="cpu")["model"]["repcodec"])
|
||||
model.quantizer.initial()
|
||||
model.eval()
|
||||
return model
|
||||
|
||||
|
||||
def load_shard(in_dir: Path, rank: int, n_shard: int) -> Tuple[np.ndarray, List[int]]:
|
||||
feat_path = in_dir / f"{rank}_{n_shard}.npy"
|
||||
len_path = in_dir / f"{rank}_{n_shard}.len"
|
||||
|
||||
with open(len_path) as fp:
|
||||
lengths = [int(line.strip()) for line in fp]
|
||||
|
||||
return np.load(feat_path.as_posix(), mmap_mode="r"), lengths
|
||||
|
||||
|
||||
def pad_data(data: List[np.ndarray]) -> List[np.ndarray]:
|
||||
max_len = max([d.shape[0] for d in data])
|
||||
data = [
|
||||
np.pad(d, [(0, max_len - d.shape[0]), (0, 0)], "constant", constant_values=0.0)
|
||||
for d in data
|
||||
]
|
||||
return data
|
||||
|
||||
|
||||
def make_batch_data(data: np.ndarray, shard_lengths: List[int], batch_size: int):
|
||||
batch_data = []
|
||||
batch_lens = []
|
||||
offsets = np.cumsum([0] + shard_lengths)
|
||||
assert len(data) == offsets[-1], f"{len(data)} {offsets[-1]}"
|
||||
|
||||
# from longest to shortest
|
||||
for i in range(len(shard_lengths)):
|
||||
if batch_size > len(batch_data):
|
||||
batch_data.append(data[offsets[i]: offsets[i + 1]])
|
||||
batch_lens.append(shard_lengths[i])
|
||||
else:
|
||||
yield {
|
||||
"data": torch.tensor(np.stack(pad_data(batch_data)), dtype=torch.float), # (bsz, seq len, hidden dim)
|
||||
"lengths": batch_lens
|
||||
}
|
||||
batch_data = [data[offsets[i]: offsets[i + 1]]]
|
||||
batch_lens = [shard_lengths[i]]
|
||||
if len(batch_data) > 0:
|
||||
yield {
|
||||
"data": torch.tensor(np.stack(pad_data(batch_data)), dtype=torch.float),
|
||||
"lengths": batch_lens
|
||||
}
|
||||
|
||||
|
||||
def tokenize_batch(model: RepCodec, batch: dict, device: str) -> List[List[int]]:
|
||||
with torch.no_grad():
|
||||
data = batch["data"].transpose(1, 2).to(device) # (bsz, hidden dim, seq len)
|
||||
x = model.encoder(data)
|
||||
z = model.projector(x)
|
||||
_, idx = model.quantizer.codebook.forward_index(z.transpose(2, 1))
|
||||
|
||||
# when bsz=1: (1, seq len)
|
||||
if idx.dim() == 2:
|
||||
return idx.cpu().data.numpy().tolist()
|
||||
# when bsz>1: (1, bsz, seq len)
|
||||
tokens = idx.cpu().data.numpy().tolist()[0]
|
||||
res = []
|
||||
batch_lens = batch["lengths"]
|
||||
for i in range(len(tokens)):
|
||||
n_tokens = batch_lens[i]
|
||||
res.append(tokens[i][:n_tokens])
|
||||
return res
|
||||
|
||||
|
||||
def load_tsv(path: str):
|
||||
with open(path) as fp:
|
||||
root = fp.readline().strip()
|
||||
names = []
|
||||
for line in fp:
|
||||
names.append(line.strip().split("\t")[0])
|
||||
return root, names
|
||||
|
||||
|
||||
def cli():
|
||||
args = parse_args()
|
||||
device = "cuda" if args.use_gpu else "cpu"
|
||||
|
||||
model = load_model(model_path=args.model, config_path=args.model_config_path)
|
||||
model.to(device)
|
||||
|
||||
in_dir = Path(args.in_dir)
|
||||
n_shard = args.n_shard
|
||||
batch_size = args.batch_size
|
||||
|
||||
root_dir, file_names = load_tsv(args.tsv_path)
|
||||
|
||||
output_dir = args.out_dir
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
processed_cnt = 0
|
||||
pbar = tqdm(total=len(file_names))
|
||||
with open(os.path.join(output_dir, "tokens"), mode="w+") as fp:
|
||||
fp.write(f"{root_dir}\n")
|
||||
|
||||
for rank in range(n_shard):
|
||||
shard_data, shard_lengths = load_shard(in_dir, rank, n_shard)
|
||||
for batch in make_batch_data(shard_data, shard_lengths, batch_size=batch_size):
|
||||
batch_tokens = tokenize_batch(model, batch, device)
|
||||
|
||||
for tokens in batch_tokens:
|
||||
fp.write(f"{file_names[processed_cnt]}\t{' '.join(map(str, tokens))}\n")
|
||||
processed_cnt += 1
|
||||
|
||||
pbar.update(len(batch_tokens))
|
||||
assert processed_cnt == len(file_names), f"# lines of tsv do not match # of representations!"
|
||||
|
||||
pbar.close()
|
||||
print("Tokenize successfully!")
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
cli()
|
||||
Reference in New Issue
Block a user