73 lines
2.3 KiB
Python
73 lines
2.3 KiB
Python
import re
|
|
import os
|
|
import torch
|
|
import torchaudio
|
|
from init import Global
|
|
from graph_rag import GraphRAGMemory
|
|
from modelscope.pipelines import pipeline
|
|
from modelscope.utils.constant import Tasks
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from speechbrain.inference import SpeakerRecognition
|
|
|
|
device = Global.device
|
|
|
|
class SpeakerVerification:
|
|
def __init__(self):
|
|
self.verification = SpeakerRecognition.from_hparams(
|
|
source="speechbrain/spkrec-ecapa-voxceleb",
|
|
run_opts={"device": device}
|
|
)
|
|
|
|
self._warmup()
|
|
|
|
def _warmup(self):
|
|
dummy_signal = torch.randn(1, 16000).to(device)
|
|
self.verification.encode_batch(dummy_signal)
|
|
|
|
def verify_speaker(self, audio_file1, audio_file2):
|
|
signal1, _ = torchaudio.load(audio_file1)
|
|
signal2, _ = torchaudio.load(audio_file2)
|
|
|
|
if device == "cuda":
|
|
signal1 = signal1.to(device)
|
|
signal2 = signal2.to(device)
|
|
|
|
score, _ = self.verification.verify_batch(signal1, signal2)
|
|
|
|
return score.item()
|
|
|
|
class SenseVoice:
|
|
def __init__(self):
|
|
self.model = pipeline(
|
|
task=Tasks.auto_speech_recognition,
|
|
model='iic/SenseVoiceSmall',
|
|
model_revision="master",
|
|
device=device
|
|
)
|
|
|
|
self._warmup()
|
|
|
|
def _warmup(self):
|
|
if os.path.exists('voice.wav'):
|
|
self.infer()
|
|
|
|
def infer(self, voice_path='voice.wav'):
|
|
result = self.model(voice_path)
|
|
pattern = r"<\|(.+?)\|><\|(.+?)\|><\|(.+?)\|><\|(.+?)\|>(.+)"
|
|
match = re.match(pattern, result[0]['text'])
|
|
if match:
|
|
language, emotion, audio_type, itn, text = match.groups()
|
|
text = f"<{emotion}>{text}"
|
|
else:
|
|
text = ''
|
|
return text
|
|
|
|
with ThreadPoolExecutor(max_workers=2) as executor:
|
|
futures = [
|
|
executor.submit(lambda: setattr(Global, 'sense_voice', SenseVoice())),
|
|
executor.submit(lambda: setattr(Global, 'speaker_verifier', SpeakerVerification())),
|
|
executor.submit(lambda: setattr(Global, 'memory', GraphRAGMemory()))
|
|
]
|
|
|
|
for future in futures:
|
|
future.result() |