Duplicate from sergeipetrov/asrdiarization-handler

Co-authored-by: Sergei Petrov <[email protected]>

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,23 @@
+ASR+Diarization handler that works natively with Inference Endpoints.
+
+Example payload:
+```python
+import base64
+import requests
+
+API_URL = "<your endpoint URL>"
+filepath = "/path/to/audio"
+
+with open(filepath, 'rb') as f:
+    audio_encoded = base64.b64encode(f.read()).decode("utf-8")
+
+data = {
+    "inputs": audio_encoded,
+    "parameters": {
+        "batch_size": 24
+    }
+}
+
+resp = requests.post(API_URL, json=data, headers={"Authorization": "Bearer <your token>"})
+print(resp.json())
+```

config.py

@@ -0,0 +1,33 @@
+import logging
+
+from pydantic import BaseModel
+from pydantic_settings import BaseSettings
+from typing import Optional, Literal
+
+logger = logging.getLogger(__name__)
+
+
+class ModelSettings(BaseSettings):
+    asr_model: str
+    assistant_model: Optional[str]
+    diarization_model: Optional[str]
+    hf_token: Optional[str]
+
+
+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
+
+
+model_settings = ModelSettings()
+
+logger.info(f"asr model: {model_settings.asr_model}")
+logger.info(f"assist model: {model_settings.assistant_model}")
+logger.info(f"diar model: {model_settings.diarization_model}")

diarization_utils.py

@@ -0,0 +1,141 @@
+import torch
+import numpy as np
+from torchaudio import functional as F
+from transformers.pipelines.audio_utils import ffmpeg_read
+from starlette.exceptions import HTTPException
+import sys
+
+# Code from insanely-fast-whisper:
+# https://github.com/Vaibhavs10/insanely-fast-whisper
+
+import logging
+logger = logging.getLogger(__name__)
+
+def preprocess_inputs(inputs, sampling_rate):
+    inputs = ffmpeg_read(inputs, sampling_rate)
+
+    if sampling_rate != 16000:
+        inputs = F.resample(
+            torch.from_numpy(inputs), sampling_rate, 16000
+        ).numpy()
+
+    if len(inputs.shape) != 1:
+        logger.error(f"Diarization pipeline expecs single channel audio, received {inputs.shape}")
+        raise HTTPException(
+            status_code=400,
+            detail=f"Diarization pipeline expecs single channel audio, received {inputs.shape}"
+        )
+
+    # diarization model expects float32 torch tensor of shape `(channels, seq_len)`
+    diarizer_inputs = torch.from_numpy(inputs).float()
+    diarizer_inputs = diarizer_inputs.unsqueeze(0)
+
+    return inputs, diarizer_inputs
+
+
+def diarize_audio(diarizer_inputs, diarization_pipeline, parameters):
+    diarization = diarization_pipeline(
+        {"waveform": diarizer_inputs, "sample_rate": parameters.sampling_rate},
+        num_speakers=parameters.num_speakers,
+        min_speakers=parameters.min_speakers,
+        max_speakers=parameters.max_speakers,
+    )
+
+    segments = []
+    for segment, track, label in diarization.itertracks(yield_label=True):
+        segments.append(
+            {
+                "segment": {"start": segment.start, "end": segment.end},
+                "track": track,
+                "label": label,
+            }
+        )
+
+    # diarizer output may contain consecutive segments from the same speaker (e.g. {(0 -> 1, speaker_1), (1 -> 1.5, speaker_1), ...})
+    # we combine these segments to give overall timestamps for each speaker's turn (e.g. {(0 -> 1.5, speaker_1), ...})
+    new_segments = []
+    prev_segment = cur_segment = segments[0]
+
+    for i in range(1, len(segments)):
+        cur_segment = segments[i]
+
+        # check if we have changed speaker ("label")
+        if cur_segment["label"] != prev_segment["label"] and i < len(segments):
+            # add the start/end times for the super-segment to the new list
+            new_segments.append(
+                {
+                    "segment": {
+                        "start": prev_segment["segment"]["start"],
+                        "end": cur_segment["segment"]["start"],
+                    },
+                    "speaker": prev_segment["label"],
+                }
+            )
+            prev_segment = segments[i]
+
+    # add the last segment(s) if there was no speaker change
+    new_segments.append(
+        {
+            "segment": {
+                "start": prev_segment["segment"]["start"],
+                "end": cur_segment["segment"]["end"],
+            },
+            "speaker": prev_segment["label"],
+        }
+    )
+
+    return new_segments
+
+
+def post_process_segments_and_transcripts(new_segments, transcript, group_by_speaker) -> list:
+    # get the end timestamps for each chunk from the ASR output
+    end_timestamps = np.array(
+        [chunk["timestamp"][-1] if chunk["timestamp"][-1] is not None else sys.float_info.max for chunk in transcript])
+    segmented_preds = []
+
+    # align the diarizer timestamps and the ASR timestamps
+    for segment in new_segments:
+        # get the diarizer end timestamp
+        end_time = segment["segment"]["end"]
+        # find the ASR end timestamp that is closest to the diarizer's end timestamp and cut the transcript to here
+        upto_idx = np.argmin(np.abs(end_timestamps - end_time))
+
+        if group_by_speaker:
+            segmented_preds.append(
+                {
+                    "speaker": segment["speaker"],
+                    "text": "".join(
+                        [chunk["text"] for chunk in transcript[: upto_idx + 1]]
+                    ),
+                    "timestamp": (
+                        transcript[0]["timestamp"][0],
+                        transcript[upto_idx]["timestamp"][1],
+                    ),
+                }
+            )
+        else:
+            for i in range(upto_idx + 1):
+                segmented_preds.append({"speaker": segment["speaker"], **transcript[i]})
+
+        # crop the transcripts and timestamp lists according to the latest timestamp (for faster argmin)
+        transcript = transcript[upto_idx + 1:]
+        end_timestamps = end_timestamps[upto_idx + 1:]
+
+        if len(end_timestamps) == 0:
+            break
+
+    return segmented_preds
+
+
+def diarize(diarization_pipeline, file, parameters, asr_outputs):
+    _, diarizer_inputs = preprocess_inputs(file, parameters.sampling_rate)
+
+    segments = diarize_audio(
+        diarizer_inputs, 
+        diarization_pipeline, 
+        parameters
+    )
+
+    return post_process_segments_and_transcripts(
+        segments, asr_outputs["chunks"], group_by_speaker=False
+    )

handler.py

@@ -0,0 +1,103 @@
+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:
+            # diarization pipeline doesn't raise if there is no token
+            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"],
+        }

requirements.txt

@@ -0,0 +1,8 @@
+accelerate==0.27.2
+torch==2.2.1
+pyannote-audio==3.1.1
+transformers==4.38.2
+numpy==1.26.4
+torchaudio==2.2.1
+pydantic==2.6.3
+pydantic-settings==2.2.1