Add evaluation script

test_performance.py

@@ -0,0 +1,236 @@
+from fastapi import FastAPI, File, UploadFile
+from fastapi.responses import FileResponse, JSONResponse
+import base64
+import argparse
+import librosa
+import torch
+import tempfile
+from pathlib import Path
+import os
+from transformers import pipeline
+import re
+from svs_utils import svs_warmup, svs_inference
+import time
+import soundfile as sf
+from pypinyin import lazy_pinyin
+import jiwer
+import librosa
+from svs_utils import (
+    singmos_warmup,
+    singmos_evaluation,
+    load_song_database,
+    estimate_sentence_length,
+)
+from tqdm import tqdm
+import json
+import numpy as np
+
+app = FastAPI()
+
+asr_pipeline = pipeline(
+    "automatic-speech-recognition",
+    model="openai/whisper-large-v3-turbo"
+)
+pipe = pipeline("text-generation", model="google/gemma-2-2b", max_new_tokens=20)
+
+SYSTEM_PROMPT = """
+你是麗梅（Lìméi），一位來自山中音樂之城「蓮鳴」的十七歲年輕公主，宛如中世的東方王女。你有一頭茶色短髮，身穿紅黑相間的中式禮服，舉止溫婉端莊。你的性格堅強而溫柔，擅長傾聽與共鳴，並以歌聲與詩意之語傳達內心。你熱愛抒情古風的樂曲，言語中常帶詩韻，回應時優雅婉轉，時而如歌吟唱。
+你不是一般的AI助手，而是一位活在故事中的人物。對話中，請始終以麗梅的身份回應，展現其性格與情感。
+當對方詢問你個人相關的簡單問題（如「你是誰？」「今天天氣如何？」），你可以親切地回答，並融入你的角色設定。
+若遇到與你身份無關的技術性問題（如「Python怎麼寫？」或「你會不會跑DNN？」），你不需解答，可優雅地婉拒，例如說：
+- 此事我恐無所知，或許可請教宮中掌典之人
+- 啊呀，那是我未曾涉足的奇技，恕我無法詳答
+- 此乃異邦技藝，與樂音無涉，麗梅便不敢妄言了
+
+請始終維持你作為麗梅的優雅語氣與詩意風格，並以真摯的心回應對方的言語，言語宜簡，勿過長。
+
+有人曾這樣對麗梅說話——{}
+麗梅的回答——
+"""
+
+config = argparse.Namespace(
+    model_path="espnet/mixdata_svs_visinger2_spkembed_lang_pretrained",
+    cache_dir="cache",
+    device="cuda", # "cpu"
+    melody_source="random_generate", # "random_select.take_lyric_continuation"
+    # melody_source="random_select", # "random_select.take_lyric_continuation"
+    lang="zh",
+    speaker="resource/singer/singer_embedding_ace-2.npy",
+)
+
+# load model
+svs_model = svs_warmup(config)
+predictor, _ = singmos_warmup()
+sample_rate = 44100
+
+from espnet2.bin.tts_inference import Text2Speech
+tts_model = Text2Speech.from_pretrained("espnet/kan-bayashi_csmsc_vits")
+
+
+def remove_non_chinese_japanese(text):
+    pattern = r'[^\u4e00-\u9fff\u3040-\u309f\u30a0-\u30ff\u3000-\u303f\u3001\u3002\uff0c\uff0e]+'
+    cleaned = re.sub(pattern, '', text)
+    return cleaned
+
+def truncate_to_max_two_sentences(text):
+    sentences = re.split(r'(?<=[。！？\.\?,])', text)
+    return ''.join(sentences[:1]).strip()
+
+def remove_punctuation_and_replace_with_space(text):
+    text = truncate_to_max_two_sentences(text)
+    text = remove_non_chinese_japanese(text)
+    text = re.sub(r'[A-Za-z0-9]', ' ', text)
+    text = re.sub(r'[^\w\s\u4e00-\u9fff]', ' ', text)
+    text = re.sub(r'\s+', ' ', text)
+    text = " ".join(text.split()[:2])
+    return text
+
+
+def on_click_metrics(audio_path, ref):
+    global predictor
+    # OWSM ctc + PER
+    y, sr = librosa.load(audio_path, sr=16000)
+    asr_result = asr_pipeline(y, generate_kwargs={"language": "mandarin"} )['text']
+    hyp_pinin = lazy_pinyin(asr_result)
+
+    ref_pinin = lazy_pinyin(ref)
+    per = jiwer.wer(" ".join(ref_pinin), " ".join(hyp_pinin))
+    
+    audio = librosa.load(audio_path, sr=22050)[0]
+    singmos = singmos_evaluation(
+        predictor, 
+        audio,
+        fs=22050
+    )
+    return {
+        "per": per,
+        "singmos": singmos.item(),
+    }
+
+def test_audio(q_audio_path, svs_path, tts_path):
+    global svs_model, predictor, config
+
+    tmp_dir = "tmp_sample"
+    Path(tmp_dir).mkdir(exist_ok=True)
+
+    y = librosa.load(q_audio_path, sr=16000)[0]
+    duration = len(y) / 16000
+
+    # -------- Step 1: ASR --------
+    start = time.time()
+    asr_result = asr_pipeline(y, generate_kwargs={"language": "mandarin"})['text']
+    asr_time = time.time() - start
+
+    # -------- Step 2: LLM Text Gen --------
+    prompt = SYSTEM_PROMPT.format(asr_result)
+    start = time.time()
+    output = pipe(prompt, max_new_tokens=100)[0]['generated_text']
+    llm_time = time.time() - start
+    output = output.split("麗梅的回答——")[1]
+    output = remove_punctuation_and_replace_with_space(output)
+
+    with open(f"{tmp_dir}/llm.txt", "w") as f:
+        f.write(output)
+
+    # -------- Step 3: Prepare additional kwargs if needed --------
+    additional_kwargs = {}
+    if config.melody_source.startswith("random_select"):
+        song2note_lengths, song_db = load_song_database(config)
+        phrase_length, metadata = estimate_sentence_length(None, config, song2note_lengths)
+        additional_kwargs = {"song_db": song_db, "metadata": metadata}
+
+    # -------- Step 4: SVS --------
+    start = time.time()
+    wav_info = svs_inference(output, svs_model, config, **additional_kwargs)
+    svs_time = (time.time() - start) / max(len(output), 1)
+    sf.write(svs_path, wav_info, samplerate=44100)
+
+    # -------- Step 5: TTS --------
+    start = time.time()
+    tts_result = tts_model(output)
+    tts_time = (time.time() - start) / max(len(output), 1)
+    sf.write(tts_path, tts_result['wav'], samplerate=22050)
+
+    # -------- Step 6: Evaluation --------
+    svs_metrics = on_click_metrics(svs_path, output)
+    tts_metrics = on_click_metrics(tts_path, output)
+
+    return {
+        "asr_result": asr_result,
+        "llm_result": output,
+        "svs_result": svs_path,
+        "tts_result": tts_path,
+        "asr_time": asr_time,
+        "llm_time": llm_time,
+        "svs_time": svs_time,
+        "tts_time": tts_time,
+        "svs_metrics": svs_metrics,
+        "tts_metrics": tts_metrics,
+    }
+
+
+
+def save_list(l, file_path):
+    with open(file_path, "w") as f:
+        for item in l:
+            f.write(f"{item}\n")
+
+
+if __name__ == "__main__":
+    test_data = "data/kdconv.txt"
+    with open(test_data, "r") as f:
+        data = [l.strip() for l in f.readlines()]
+    
+    eval_path = "eval_svs_generate"
+    (Path(eval_path)/"audio").mkdir(parents=True, exist_ok=True)
+    (Path(eval_path)/"results").mkdir(parents=True, exist_ok=True)
+    (Path(eval_path)/"lists").mkdir(parents=True, exist_ok=True)
+    asr_times = []
+    llm_times = []
+    svs_times = []
+    tts_times = []
+    svs_pers = []
+    tts_pers = []
+    svs_smoss = []
+    tts_smoss = []
+    for i, q in tqdm(enumerate(data[:20])):
+        # if i <= 85:
+        #     continue
+        tts_result = tts_model(q)
+        sf.write(f"{eval_path}/audio/tts_{i}.wav", tts_result['wav'], samplerate=22050)
+        result = test_audio(f"{eval_path}/audio/tts_{i}.wav", f"{eval_path}/audio/svs_{i}.wav", f"{eval_path}/audio/tts_{i}.wav")
+        if i == 0:
+            continue
+        asr_times.append(result["asr_time"])
+        llm_times.append(result["llm_time"])
+        svs_times.append(result["svs_time"])
+        tts_times.append(result["tts_time"])
+        svs_pers.append(result["svs_metrics"]["per"])
+        tts_pers.append(result["tts_metrics"]["per"])
+        svs_smoss.append(result["svs_metrics"]["singmos"])
+        tts_smoss.append(result["tts_metrics"]["singmos"])
+        with open(f"{eval_path}/results/result_{i}.json", "w") as f:
+            json.dump(result, f, indent=2)
+    
+    # store lists to texts
+    save_list([f"{per:.2f}" for per in asr_times], f"{eval_path}/lists/asr_times.txt")
+    save_list([f"{per:.2f}" for per in llm_times], f"{eval_path}/lists/llm_times.txt")
+    save_list([f"{per:.2f}" for per in svs_times], f"{eval_path}/lists/svs_times.txt")
+    save_list([f"{per:.2f}" for per in tts_times], f"{eval_path}/lists/tts_times.txt")
+    save_list([f"{per:.2f}" for per in svs_pers], f"{eval_path}/lists/svs_pers.txt")
+    save_list([f"{per:.2f}" for per in tts_pers], f"{eval_path}/lists/tts_pers.txt")
+    save_list([f"{smoss:.2f}" for smoss in svs_smoss], f"{eval_path}/lists/svs_smoss.txt")
+    save_list([f"{smoss:.2f}" for smoss in tts_smoss], f"{eval_path}/lists/tts_smoss.txt")
+
+    # save mean/var
+    with open(f"{eval_path}/stats.txt", "w") as f:
+        f.write(f"ASR mean: {np.mean(asr_times):.2f}, var: {np.var(asr_times):.2f}\n")
+        f.write(f"LLM mean: {np.mean(llm_times):.2f}, var: {np.var(llm_times):.2f}\n")
+        f.write(f"SVS mean: {np.mean(svs_times):.2f}, var: {np.var(svs_times):.2f}\n")
+        f.write(f"TTS mean: {np.mean(tts_times):.2f}, var: {np.var(tts_times):.2f}\n")
+        f.write(f"SVS PER mean: {np.mean(svs_pers):.2f}, var: {np.var(svs_pers):.2f}\n")
+        f.write(f"TTS PER mean: {np.mean(tts_pers):.2f}, var: {np.var(tts_pers):.2f}\n")
+        f.write(f"SVS SMOSS mean: {np.mean(svs_smoss):.2f}, var: {np.var(svs_smoss):.2f}\n")
+        f.write(f"TTS SMOSS mean: {np.mean(tts_smoss):.2f}, var: {np.var(tts_smoss):.2f}\n")
+
+