Revert "change g2p and other fix" (#40)

text/japanese_bert.py

@@ -2,12 +2,7 @@ import torch
 from transformers import AutoTokenizer, AutoModelForMaskedLM
 import sys
 
-BERT = "./bert/bert-large-japanese-v2"
-
-tokenizer = AutoTokenizer.from_pretrained(BERT)
-# bert-large model has 25 hidden layers.You can decide which layer to use by setting this variable to a specific value
-# default value is 3(untested)
-BERT_LAYER = 3
+tokenizer = AutoTokenizer.from_pretrained("./bert/bert-base-japanese-v3")
 
 models = dict()
 
@@ -29,13 +24,13 @@ def get_bert_feature(text, word2ph, device=None):
         inputs = tokenizer(text, return_tensors="pt")
         for i in inputs:
             inputs[i] = inputs[i].to(device)
-        res = model(**inputs, output_hidden_states=True)
-        res = res["hidden_states"][BERT_LAYER]
+        res = models[device](**inputs, output_hidden_states=True)
+        res = torch.cat(res["hidden_states"][-3:-2], -1)[0].cpu()
     assert inputs["input_ids"].shape[-1] == len(word2ph)
     word2phone = word2ph
     phone_level_feature = []
     for i in range(len(word2phone)):
-        repeat_feature = res[0][i].repeat(word2phone[i], 1)
+        repeat_feature = res[i].repeat(word2phone[i], 1)
         phone_level_feature.append(repeat_feature)
 
     phone_level_feature = torch.cat(phone_level_feature, dim=0)