From cf6498ef3055127684b3ca143fd35412d4c32cb7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stardust=C2=B7=E5=87=8F?= <star_dust_chen@foxmail.com>
Date: Sun, 8 Oct 2023 22:49:36 +0800
Subject: [PATCH] =?UTF-8?q?=E5=88=9D=E6=AD=A5=EF=BC=8C=E7=9D=A1=E8=A7=89?=
 =?UTF-8?q?=E6=98=8E=E5=A4=A9=E7=BB=A7=E7=BB=AD=E5=86=99=EF=BC=88?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 add_punc.py | 82 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 82 insertions(+)
 create mode 100644 add_punc.py

diff --git a/add_punc.py b/add_punc.py
new file mode 100644
index 0000000..c299de1
--- /dev/null
+++ b/add_punc.py
@@ -0,0 +1,82 @@
+import json
+import os
+import multiprocessing
+from modelscope.pipelines import pipeline
+from modelscope.utils.constant import Tasks
+from tqdm import tqdm
+import logging
+
+logging.getLogger("modelscope").setLevel(logging.ERROR)
+
+metadata_path = "filelists/genshin.list"
+output_path = "filelists/genshin_punc.list"
+
+os.makedirs("tmp", exist_ok=True)
+
+
+def split_list(lst, n):
+    avg = len(lst) // n  # 每份的平均长度
+    rem = len(lst) % n  # 剩余的长度
+
+    result = []
+    start = 0
+
+    for i in range(n):
+        end = start + avg + (1 if i < rem else 0)  # 计算当前份的结束位置
+        result.append(lst[start:end])  # 切片并将结果添加到结果列表中
+        start = end  # 更新下一份的起始位置
+
+    return result
+
+
+def process(audios, pid):
+    print("loading model")
+    inference_pipeline = pipeline(
+        task=Tasks.punctuation,
+        model='damo/punc_ct-transformer_zh-cn-common-vocab272727-pytorch'
+    )
+    print("loaded model")
+
+    total_cnt = 0
+    skip_cnt = 0
+    for audiofile in tqdm(audios):
+        for k, sentence in enumerate(audiofile['segments']):
+            confidence = sentence['confidence']
+            if confidence >= 0.95:
+                text = sentence['text']
+                rec_result = inference_pipeline(text_in=text)
+                sentence['text'] = rec_result['text']
+                total_cnt += 1
+            else:
+                skip_cnt += 1
+        print("preprocessed text count :", total_cnt, "skip count :", skip_cnt)
+    with open(f"tmp/{pid}.json", 'w') as f:
+        json.dump(audios, f, ensure_ascii=False, indent=2)
+
+
+def process_wrapper(args):
+    audios, pid = args
+    process(audios, pid)
+
+
+if __name__ == '__main__':
+    # Define the number of processes
+    n_process = 8
+
+    # Split the audios into chunks for multiprocessing
+    with open(metadata_path, "r") as f:
+        lines = f.readlines()
+    audio_chunks = split_list(audios, n_process)
+
+    # Create a pool of worker processes
+    pool = multiprocessing.Pool(processes=n_process)
+
+    # Map the process function to each chunk of audios
+    pool.map(process_wrapper, zip(audio_chunks, range(n_process)))
+
+    # Close the pool and wait for the work to finish
+    pool.close()
+    pool.join()
+
+    # Merge output files
+    merged_data = []