Spaces:
Sleeping
Sleeping
File size: 2,732 Bytes
9712d04 780954b 9712d04 |
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 |
import os
import tempfile
from typing import Optional
import numpy as np
import soundfile as sf
try:
from funasr import AutoModel
except ImportError:
AutoModel = None
from .base import AbstractASRModel
from .registry import register_asr_model
@register_asr_model("funasr/paraformer-zh")
class ParaformerASR(AbstractASRModel):
def __init__(
self, model_id: str, device: str = "auto", cache_dir: str = "cache", **kwargs
):
super().__init__(model_id, device, cache_dir, **kwargs)
if AutoModel is None:
raise ImportError(
"funasr is not installed. Please install it with: pip3 install -U funasr"
)
model_name = model_id.replace("funasr/", "")
language = model_name.split("-")[1]
if language == "zh":
self.language = "mandarin"
elif language == "en":
self.language = "english"
else:
raise ValueError(
f"Language cannot be determined. {model_id} is not supported"
)
try:
original_cache_dir = os.getenv("MODELSCOPE_CACHE")
os.makedirs(cache_dir, exist_ok=True)
os.environ["MODELSCOPE_CACHE"] = cache_dir
self.model = AutoModel(
model=model_name,
model_revision="v2.0.4",
vad_model="fsmn-vad",
vad_model_revision="v2.0.4",
punc_model="ct-punc-c",
punc_model_revision="v2.0.4",
device=device,
)
if original_cache_dir:
os.environ["MODELSCOPE_CACHE"] = original_cache_dir
else:
del os.environ["MODELSCOPE_CACHE"]
except Exception as e:
raise ValueError(f"Error loading Paraformer model: {e}")
def transcribe(
self,
audio: np.ndarray,
audio_sample_rate: int,
language: Optional[str] = None,
**kwargs,
) -> str:
if language and language != self.language:
raise ValueError(
f"Paraformer model {self.model_id} only supports {self.language} language, but {language} was requested"
)
try:
with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
sf.write(f.name, audio, audio_sample_rate)
temp_file = f.name
result = self.model.generate(input=temp_file, batch_size_s=300, **kwargs)
os.unlink(temp_file)
print(f"Transcription result: {result}, type: {type(result)}")
return result[0]["text"]
except Exception as e:
raise ValueError(f"Error during transcription: {e}")
|