|
import logging |
|
import torch |
|
import os |
|
import base64 |
|
|
|
from pyannote.audio import Pipeline |
|
from transformers import pipeline, AutoModelForCausalLM |
|
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.asr_pipeline = pipeline( |
|
"automatic-speech-recognition", |
|
model=model_settings.asr_model, |
|
torch_dtype=torch_dtype, |
|
device=device |
|
) |
|
|
|
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 |
|
} |
|
|
|
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 during ASR inference: {str(e)}") |
|
raise HTTPException(status_code=500, detail=f"Unknown error during ASR inference: {str(e)}") |
|
|
|
return { |
|
"chunks": asr_outputs["chunks"], |
|
"text": asr_outputs["text"], |
|
} |