freedakgmail
6116b1f3c6
Update API Documentation / build-api-docs (push) Has been cancelled
Add complete FunASR codebase including models, runtime, and documentation.
50 lines
1.5 KiB
Python
Executable File
50 lines
1.5 KiB
Python
Executable File
from transformers import AutoTokenizer, AutoModel, pipeline
|
|
import numpy as np
|
|
import sys
|
|
import os
|
|
import torch
|
|
from kaldiio import WriteHelper
|
|
import re
|
|
|
|
text_file_json = sys.argv[1]
|
|
out_ark = sys.argv[2]
|
|
out_scp = sys.argv[3]
|
|
out_shape = sys.argv[4]
|
|
device = int(sys.argv[5])
|
|
model_path = sys.argv[6]
|
|
|
|
model = AutoModel.from_pretrained(model_path)
|
|
tokenizer = AutoTokenizer.from_pretrained(model_path)
|
|
extractor = pipeline(task="feature-extraction", model=model, tokenizer=tokenizer, device=device)
|
|
|
|
with open(text_file_json, "r") as f:
|
|
js = f.readlines()
|
|
|
|
|
|
f_shape = open(out_shape, "w")
|
|
with WriteHelper("ark,scp:{},{}".format(out_ark, out_scp)) as writer:
|
|
with torch.no_grad():
|
|
for idx, line in enumerate(js):
|
|
id, tokens = line.strip().split(" ", 1)
|
|
tokens = re.sub(" ", "", tokens.strip())
|
|
tokens = " ".join([j for j in tokens])
|
|
token_num = len(tokens.split(" "))
|
|
outputs = extractor(tokens)
|
|
outputs = np.array(outputs)
|
|
embeds = outputs[0, 1:-1, :]
|
|
|
|
token_num_embeds, dim = embeds.shape
|
|
if token_num == token_num_embeds:
|
|
writer(id, embeds)
|
|
shape_line = "{} {},{}\n".format(id, token_num_embeds, dim)
|
|
f_shape.write(shape_line)
|
|
else:
|
|
print(
|
|
"{}, size has changed, {}, {}, {}".format(
|
|
id, token_num, token_num_embeds, tokens
|
|
)
|
|
)
|
|
|
|
|
|
f_shape.close()
|