forked from espnet/espnet
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathjson2trn_mt.py
More file actions
executable file
·124 lines (103 loc) · 3.84 KB
/
json2trn_mt.py
File metadata and controls
executable file
·124 lines (103 loc) · 3.84 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
#!/usr/bin/env python3
# encoding: utf-8
# Copyright 2018 Kyoto University (Hirofumi Inaguma)
# Apache 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
# NOTE: this is made for machine translation
import argparse
import codecs
import json
import logging
import sys
from espnet.utils.cli_utils import get_commandline_args
def get_parser():
parser = argparse.ArgumentParser(
description="convert json to machine translation transcription",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
parser.add_argument("json", type=str, help="json files")
parser.add_argument("dict", type=str, help="dict for target language")
parser.add_argument("--refs", type=str, nargs="+", help="ref for all speakers")
parser.add_argument("--hyps", type=str, nargs="+", help="hyp for all outputs")
parser.add_argument("--srcs", type=str, nargs="+", help="src for all outputs")
parser.add_argument(
"--dict-src",
type=str,
help="dict for source language",
default=False,
nargs="?",
)
return parser
def main(args):
parser = get_parser()
args = parser.parse_args(args)
convert(args.json, args.dict, args.refs, args.hyps, args.srcs, args.dict_src)
def convert(jsonf, dic, refs, hyps, srcs, dic_src):
# logging info
logfmt = "%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s"
logging.basicConfig(level=logging.INFO, format=logfmt)
logging.info(get_commandline_args())
logging.info("reading %s", jsonf)
with codecs.open(jsonf, "r", encoding="utf-8") as f:
j = json.load(f)
# target dictionary
logging.info("reading %s", dic)
with codecs.open(dic, "r", encoding="utf-8") as f:
dictionary = f.readlines()
char_list_tgt = [entry.split(" ")[0] for entry in dictionary]
char_list_tgt.insert(0, "<blank>")
char_list_tgt.append("<eos>")
# source dictionary
logging.info("reading %s", dic_src)
if dic_src:
with codecs.open(dic_src, "r", encoding="utf-8") as f:
dictionary = f.readlines()
char_list_src = [entry.split(" ")[0] for entry in dictionary]
char_list_src.insert(0, "<blank>")
char_list_src.append("<eos>")
if hyps:
hyp_file = codecs.open(hyps[0], "w", encoding="utf-8")
ref_file = codecs.open(refs[0], "w", encoding="utf-8")
if srcs:
src_file = codecs.open(srcs[0], "w", encoding="utf-8")
for x in j["utts"]:
# hyps
if hyps:
seq = [
char_list_tgt[int(i)]
for i in j["utts"][x]["output"][0]["rec_tokenid"].split()
]
hyp_file.write(" ".join(seq).replace("<eos>", "")),
hyp_file.write(
" (" + j["utts"][x]["utt2spk"].replace("-", "_") + "-" + x + ")\n"
)
# ref
seq = [
char_list_tgt[int(i)] for i in j["utts"][x]["output"][0]["tokenid"].split()
]
ref_file.write(" ".join(seq).replace("<eos>", "")),
ref_file.write(
" (" + j["utts"][x]["utt2spk"].replace("-", "_") + "-" + x + ")\n"
)
# src
if "tokenid_src" in j["utts"][x]["output"][0].keys():
if dic_src:
seq = [
char_list_src[int(i)]
for i in j["utts"][x]["output"][0]["tokenid_src"].split()
]
else:
seq = [
char_list_tgt[int(i)]
for i in j["utts"][x]["output"][0]["tokenid_src"].split()
]
src_file.write(" ".join(seq).replace("<eos>", "")),
src_file.write(
" (" + j["utts"][x]["utt2spk"].replace("-", "_") + "-" + x + ")\n"
)
if hyps:
hyp_file.close()
ref_file.close()
if srcs:
src_file.close()
if __name__ == "__main__":
main(sys.argv[1:])