wehub-resource-sync
84 lines
2.9 KiB
Python
Executable File
84 lines
2.9 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
# -*- encoding: utf-8 -*-
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
from typing import List, Tuple
|
|
|
|
from omegaconf import OmegaConf
|
|
|
|
from funasr import AutoModel
|
|
|
|
|
|
def load_jsonl(jsonl_path: str) -> Tuple[List[str], List[str]]:
|
|
keys = []
|
|
targets = []
|
|
with open(jsonl_path, "r", encoding="utf-8") as f:
|
|
for line in f:
|
|
if not line.strip():
|
|
continue
|
|
record = json.loads(line)
|
|
key = record.get("key")
|
|
if key is None and isinstance(record.get("source"), dict):
|
|
key = record["source"].get("key")
|
|
keys.append(key or "")
|
|
targets.append(record.get("target", ""))
|
|
return keys, targets
|
|
|
|
|
|
def build_model(args: argparse.Namespace):
|
|
kwargs = {}
|
|
if args.config_path and args.config_name:
|
|
cfg_path = os.path.join(args.config_path, args.config_name)
|
|
cfg = OmegaConf.load(cfg_path)
|
|
kwargs.update(OmegaConf.to_container(cfg, resolve=True))
|
|
if args.model:
|
|
kwargs["model"] = args.model
|
|
if args.init_param:
|
|
kwargs["init_param"] = args.init_param
|
|
kwargs["device"] = args.device
|
|
if args.batch_size:
|
|
kwargs["batch_size"] = args.batch_size
|
|
return AutoModel(**kwargs)
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--model", type=str, default=None, help="model name or model dir")
|
|
parser.add_argument("--config-path", type=str, default=None, help="config directory")
|
|
parser.add_argument("--config-name", type=str, default=None, help="config filename")
|
|
parser.add_argument("--init-param", type=str, default=None, help="model checkpoint path")
|
|
parser.add_argument("--input-jsonl", type=str, required=True, help="input jsonl with source/target")
|
|
parser.add_argument("--output-dir", type=str, required=True, help="output directory")
|
|
parser.add_argument("--device", type=str, default="cuda:0", help="cuda:0 or cpu")
|
|
parser.add_argument("--batch-size", type=int, default=1, help="batch size for inference")
|
|
args = parser.parse_args()
|
|
|
|
os.makedirs(args.output_dir, exist_ok=True)
|
|
|
|
keys, targets = load_jsonl(args.input_jsonl)
|
|
|
|
model = build_model(args)
|
|
results = model.generate(input=args.input_jsonl, batch_size=args.batch_size)
|
|
|
|
hyp_path = os.path.join(args.output_dir, "text.hyp")
|
|
ref_path = os.path.join(args.output_dir, "text.ref")
|
|
|
|
with open(hyp_path, "w", encoding="utf-8") as hyp_f, open(
|
|
ref_path, "w", encoding="utf-8"
|
|
) as ref_f:
|
|
for idx, result in enumerate(results):
|
|
key = keys[idx] if idx < len(keys) else result.get("key", f"utt_{idx}")
|
|
hyp = result.get("text", "")
|
|
ref = targets[idx] if idx < len(targets) else ""
|
|
hyp_f.write(f"{key} {hyp}\n")
|
|
ref_f.write(f"{key} {ref}\n")
|
|
|
|
print(f"hyp saved to: {hyp_path}")
|
|
print(f"ref saved to: {ref_path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|