File size: 4,034 Bytes
c7d2a58 6b7e6e6 c7d2a58 |
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 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 |
import logging
import torch
import os
from pyannote.audio import Pipeline
from transformers import pipeline, AutoModelForCausalLM
from diarization_utils import diarize
from huggingface_hub import HfApi
from transformers.pipelines.audio_utils import ffmpeg_read
from pydantic import Json, BaseModel, ValidationError
logger = logging.getLogger(__name__)
class InferenceConfig(BaseModel):
task: Literal["transcribe", "translate"] = "transcribe"
batch_size: int = 24
assisted: bool = False
chunk_length_s: int = 30
sampling_rate: int = 16000
language: Optional[str] = None
num_speakers: Optional[int] = None
min_speakers: Optional[int] = None
max_speakers: Optional[int] = None
class EndpointHandler():
def __init__(self):
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
logger.info(f"Using device: {device.type}")
torch_dtype = torch.float32 if device.type == "cpu" else torch.float16
self.assistant_model = AutoModelForCausalLM.from_pretrained(
os.getenv("ASSISTANT_MODEL"),
torch_dtype=torch_dtype,
low_cpu_mem_usage=True,
use_safetensors=True
) if os.getenv("ASSISTANT_MODEL") else None
if self.assistant_model:
self.assistant_model.to(device)
self.asr_pipeline = pipeline(
"automatic-speech-recognition",
model=os.getenv("ASR_MODEL"),
torch_dtype=torch_dtype,
device=device
)
if os.getenv("DIARIZATION_MODEL"):
# diarization pipeline doesn't raise if there is no token
HfApi().whoami(model_settings.hf_token)
self.diarization_pipeline = Pipeline.from_pretrained(
checkpoint_path=os.getenv("DIARIZATION_MODEL"),
use_auth_token=os.getenv("HF_TOKEN"),
)
self.diarization_pipeline.to(device)
else:
self.diarization_pipeline = None
async def __call__(self, file, parameters):
try:
parameters = InferenceConfig(**parameters)
except ValidationError as e:
logger.error(f"Error validating parameters: {e}")
raise ValidationError(f"Error validating parameters: {e}")
logger.info(f"inference parameters: {parameters}")
generate_kwargs = {
"task": parameters.task,
"language": parameters.language,
"assistant_model": self.assistant_model if parameters.assisted else None
}
try:
asr_outputs = self.asr_pipeline(
file,
chunk_length_s=parameters.chunk_length_s,
batch_size=parameters.batch_size,
generate_kwargs=generate_kwargs,
return_timestamps=True,
)
except RuntimeError as e:
logger.error(f"ASR inference error: {str(e)}")
raise HTTPException(status_code=400, detail=f"ASR inference error: {str(e)}")
except Exception as e:
logger.error(f"Unknown error diring ASR inference: {str(e)}")
raise HTTPException(status_code=500, detail=f"Unknown error diring ASR inference: {str(e)}")
if self.diarization_pipeline:
try:
transcript = diarize(self.diarization_pipeline, file, parameters, asr_outputs)
except RuntimeError as e:
logger.error(f"Diarization inference error: {str(e)}")
raise HTTPException(status_code=400, detail=f"Diarization inference error: {str(e)}")
except Exception as e:
logger.error(f"Unknown error during diarization: {str(e)}")
raise HTTPException(status_code=500, detail=f"Unknown error during diarization: {str(e)}")
else:
transcript = []
return {
"speakers": transcript,
"chunks": asr_outputs["chunks"],
"text": asr_outputs["text"],
} |