fac4f9a8ab
"text_processing" is clearer, but the import statement is longer. "nlp" is shorter and makes it clear that it is natural language processing.
81 lines
2.8 KiB
Python
81 lines
2.8 KiB
Python
import sys
|
|
from typing import Optional
|
|
|
|
import torch
|
|
from transformers import PreTrainedModel
|
|
|
|
from style_bert_vits2.constants import Languages
|
|
from style_bert_vits2.nlp import bert_models
|
|
|
|
|
|
__models: dict[torch.device | str, PreTrainedModel] = {}
|
|
|
|
|
|
def extract_bert_feature(
|
|
text: str,
|
|
word2ph: list[int],
|
|
device: torch.device | str,
|
|
assist_text: Optional[str] = None,
|
|
assist_text_weight: float = 0.7,
|
|
) -> torch.Tensor:
|
|
"""
|
|
英語のテキストから BERT の特徴量を抽出する
|
|
|
|
Args:
|
|
text (str): 英語のテキスト
|
|
word2ph (list[int]): 元のテキストの各文字に音素が何個割り当てられるかを表すリスト
|
|
device (torch.device | str): 推論に利用するデバイス
|
|
assist_text (Optional[str], optional): 補助テキスト (デフォルト: None)
|
|
assist_text_weight (float, optional): 補助テキストの重み (デフォルト: 0.7)
|
|
|
|
Returns:
|
|
torch.Tensor: BERT の特徴量
|
|
"""
|
|
|
|
if (
|
|
sys.platform == "darwin"
|
|
and torch.backends.mps.is_available()
|
|
and device == "cpu"
|
|
):
|
|
device = "mps"
|
|
if not device:
|
|
device = "cuda"
|
|
if device == "cuda" and not torch.cuda.is_available():
|
|
device = "cpu"
|
|
if device not in __models.keys():
|
|
__models[device] = bert_models.load_model(Languages.EN).to(device) # type: ignore
|
|
|
|
style_res_mean = None
|
|
with torch.no_grad():
|
|
tokenizer = bert_models.load_tokenizer(Languages.EN)
|
|
inputs = tokenizer(text, return_tensors="pt")
|
|
for i in inputs:
|
|
inputs[i] = inputs[i].to(device) # type: ignore
|
|
res = __models[device](**inputs, output_hidden_states=True)
|
|
res = torch.cat(res["hidden_states"][-3:-2], -1)[0].cpu()
|
|
if assist_text:
|
|
style_inputs = tokenizer(assist_text, return_tensors="pt")
|
|
for i in style_inputs:
|
|
style_inputs[i] = style_inputs[i].to(device) # type: ignore
|
|
style_res = __models[device](**style_inputs, output_hidden_states=True)
|
|
style_res = torch.cat(style_res["hidden_states"][-3:-2], -1)[0].cpu()
|
|
style_res_mean = style_res.mean(0)
|
|
|
|
assert len(word2ph) == res.shape[0], (text, res.shape[0], len(word2ph))
|
|
word2phone = word2ph
|
|
phone_level_feature = []
|
|
for i in range(len(word2phone)):
|
|
if assist_text:
|
|
assert style_res_mean is not None
|
|
repeat_feature = (
|
|
res[i].repeat(word2phone[i], 1) * (1 - assist_text_weight)
|
|
+ style_res_mean.repeat(word2phone[i], 1) * assist_text_weight
|
|
)
|
|
else:
|
|
repeat_feature = res[i].repeat(word2phone[i], 1)
|
|
phone_level_feature.append(repeat_feature)
|
|
|
|
phone_level_feature = torch.cat(phone_level_feature, dim=0)
|
|
|
|
return phone_level_feature.T
|