| from typing import Dict, Any | |
| from transformers import AutoModelForSpeechSeq2Seq, AutoTokenizer, pipeline | |
| class EndpointHandler(): | |
| def __init__(self, path=""): | |
| model_id = "openai/whisper-large-v3" | |
| task = "automatic-speech-recognition" | |
| self.model = AutoModelForSpeechSeq2Seq.from_pretrained(model_id) | |
| self.tokenizer = AutoTokenizer.from_pretrained(model_id) | |
| self.pipeline = pipeline(task, model=self.model, tokenizer=self.tokenizer) | |
| def __call__(self, data: Dict[str, bytes]) -> Dict[str, str]: | |
| inputs = data.pop("inputs", data) | |
| parameters = data.pop("parameters", None) | |
| if parameters is not None: | |
| result = self.pipeline(inputs, return_timestamps=True, **parameters) | |
| else: | |
| result = self.pipeline(inputs, return_timestamps=True) | |
| return {"chunks": result["chunks"]} |