|
import logging |
|
import torch |
|
import os |
|
import base64 |
|
|
|
from pyannote.audio import Pipeline |
|
from transformers import pipeline, AutoModelForCausalLM |
|
from diarization_utils import diarize |
|
from huggingface_hub import HfApi |
|
from pydantic import ValidationError |
|
from starlette.exceptions import HTTPException |
|
|
|
from config import model_settings, InferenceConfig |
|
|
|
logger = logging.getLogger(__name__) |
|
|
|
|
|
class EndpointHandler(): |
|
def __init__(self, path=""): |
|
|
|
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( |
|
model_settings.assistant_model, |
|
torch_dtype=torch_dtype, |
|
low_cpu_mem_usage=True, |
|
use_safetensors=True |
|
) if model_settings.assistant_model else None |
|
|
|
if self.assistant_model: |
|
self.assistant_model.to(device) |
|
|
|
self.asr_pipeline = pipeline( |
|
"automatic-speech-recognition", |
|
model=model_settings.asr_model, |
|
torch_dtype=torch_dtype, |
|
device=device |
|
) |
|
|
|
if model_settings.diarization_model: |
|
|
|
HfApi().whoami(model_settings.hf_token) |
|
self.diarization_pipeline = Pipeline.from_pretrained( |
|
checkpoint_path=model_settings.diarization_model, |
|
use_auth_token=model_settings.hf_token, |
|
) |
|
self.diarization_pipeline.to(device) |
|
else: |
|
self.diarization_pipeline = None |
|
|
|
|
|
def __call__(self, inputs): |
|
file = inputs.pop("inputs") |
|
file = base64.b64decode(file) |
|
parameters = inputs.pop("parameters", {}) |
|
try: |
|
parameters = InferenceConfig(**parameters) |
|
except ValidationError as e: |
|
logger.error(f"Error validating parameters: {e}") |
|
raise HTTPException(status_code=400, detail=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"], |
|
} |
