Upload folder using huggingface_hub

README.md

@@ -0,0 +1,16 @@
+---
+language: en
+tags:
+- asr
+- diarization
+pipeline_tag: automatic-speech-recognition
+---
+# ASR + Diarization Pipeline
+
+This package provides an **Automatic Speech Recognition (ASR) + Speaker Diarization** pipeline using:
+- [OpenAI Whisper](https://huggingface.co/openai/whisper-medium)
+- [Pyannote diarization](https://huggingface.co/pyannote/speaker-diarization-3.1)
+
+## Install
+```bash
+pip install git+https://huggingface.co/Capstone04/asr-diarization-pipeline

asr_diarization/__init__.py

@@ -0,0 +1 @@
+from .pipeline import ASR_Diarization

asr_diarization/inference.py

@@ -0,0 +1,13 @@
+import os
+from .pipeline import ASR_Diarization
+
+HF_TOKEN = os.environ.get("HF_TOKEN", None)
+pipe = ASR_Diarization(HF_TOKEN)
+
+def inference(inputs):
+    return pipe(inputs)
+
+def inference_with_eval(inputs, output_dir, base_name, ref_rttm=None, ref_json=None):
+    result = pipe(inputs)
+    pipe.evaluate(output_dir, base_name, ref_rttm, ref_json)
+    return result

asr_diarization/pipeline.py

@@ -0,0 +1,160 @@
+import os
+import json
+import torch
+import torchaudio
+import noisereduce as nr
+from pyannote.audio import Pipeline
+from transformers import pipeline as hf_pipeline
+import tempfile
+
+from pyannote.core import Annotation, Segment
+from pyannote.metrics.diarization import DiarizationErrorRate
+from jiwer import wer, Compose, ToLowerCase, RemovePunctuation, RemoveMultipleSpaces, Strip
+
+
+class ASR_Diarization:
+    def __init__(self, HF_TOKEN,
+                 diar_model="pyannote/speaker-diarization-3.1",
+                 asr_model="openai/whisper-medium"):
+        self.HF_TOKEN = HF_TOKEN
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        # Load diarization model
+        self.diar_pipeline = Pipeline.from_pretrained(diar_model, use_auth_token=HF_TOKEN)
+
+        # Load ASR model with timestamps
+        self.asr_pipeline = hf_pipeline(
+            "automatic-speech-recognition",
+            model=asr_model,
+            device=0 if self.device == "cuda" else -1,
+            return_timestamps=True
+        )
+
+    def run_diarization(self, audio_path):
+        diarization = self.diar_pipeline(audio_path)
+        return [
+            {"start": t.start, "end": t.end, "speaker": spk}
+            for t, _, spk in diarization.itertracks(yield_label=True)
+        ]
+
+    def run_transcription(self, audio_path, diar_json):
+        audio, sr = torchaudio.load(audio_path)
+        merged_segments = []
+        speaker_segments = {}
+
+        for seg in diar_json:
+            start, end, spk = seg["start"], seg["end"], seg["speaker"]
+            start_sample, end_sample = int(start * sr), int(end * sr)
+            chunk = audio[0, start_sample:end_sample].numpy()
+
+            reduced = nr.reduce_noise(y=chunk, sr=sr)
+            result = self.asr_pipeline(reduced)
+
+            tokens = []
+            if "chunks" in result:
+                for word_info in result["chunks"]:
+                    start_ts, end_ts = word_info.get("timestamp", (None, None)) or (None, None)
+                    tokens.append({
+                        "start": start_ts,
+                        "end": end_ts,
+                        "text": word_info["text"],
+                        "tag": "w"
+                    })
+
+            seg_dict = {
+                "speaker": spk,
+                "start": start,
+                "end": end,
+                "tokens": tokens
+            }
+            merged_segments.append(seg_dict)
+
+            if spk not in speaker_segments:
+                speaker_segments[spk] = []
+            speaker_segments[spk].append(seg_dict)
+
+        return merged_segments, list(speaker_segments.keys())
+
+    def run_pipeline(self, audio_path, output_dir=None, base_name=None,
+                     ref_rttm=None, ref_json=None):
+        diar_json = self.run_diarization(audio_path)
+        merged_segments, speakers = self.run_transcription(audio_path, diar_json)
+
+        if output_dir and base_name:
+            os.makedirs(output_dir, exist_ok=True)
+
+            # Save RTTM
+            rttm_path = os.path.join(output_dir, f"{base_name}.rttm")
+            with open(rttm_path, "w") as f:
+                for seg in diar_json:
+                    f.write(
+                        f"SPEAKER {base_name} 1 {seg['start']:.6f} "
+                        f"{seg['end']-seg['start']:.6f} <NA> <NA> "
+                        f"{seg['speaker']} <NA>\n"
+                    )
+
+            # Save transcription
+            merged_path = os.path.join(output_dir, f"{base_name}_merged_transcription.json")
+            with open(merged_path, "w") as f:
+                json.dump(merged_segments, f, indent=2)
+
+        # --- evaluation if refs are provided ---
+        eval_results = None
+        if ref_rttm or ref_json:
+            eval_results = self.evaluate(output_dir, base_name,
+                                         ref_rttm=ref_rttm, ref_json=ref_json)
+
+        return {
+            "speakers": speakers,
+            "segments": merged_segments,
+            "evaluation": eval_results
+        }
+
+    def evaluate(self, output_dir, base_name, ref_rttm=None, ref_json=None):
+        results = {}
+
+        hyp_rttm = os.path.join(output_dir, f"{base_name}.rttm")
+        hyp_json = os.path.join(output_dir, f"{base_name}_merged_transcription.json")
+
+        if ref_rttm:
+            def load_rttm(path):
+                ann = Annotation()
+                for line in open(path):
+                    if line.startswith("SPEAKER"):
+                        p = line.split()
+                        start, dur, spk = float(p[3]), float(p[4]), p[7]
+                        ann[Segment(start, start+dur)] = spk
+                return ann
+
+            der_score = DiarizationErrorRate()(load_rttm(ref_rttm), load_rttm(hyp_rttm))
+            results["DER"] = round(der_score * 100, 2)
+
+        if ref_json:
+            def load_words(path):
+                data = json.load(open(path))
+                return " ".join([tok["text"] for seg in data for tok in seg["tokens"]])
+
+            ref_text, hyp_text = load_words(ref_json), load_words(hyp_json)
+            transform = Compose([ToLowerCase(), RemovePunctuation(),
+                                 RemoveMultipleSpaces(), Strip()])
+            results["WER_raw"] = round(wer(ref_text, hyp_text), 4)
+            results["WER_normalized"] = round(wer(transform(ref_text), transform(hyp_text)), 4)
+
+        return results if results else None
+
+    def __call__(self, inputs):
+        if isinstance(inputs, dict):
+            if "audio_bytes" in inputs:
+                audio_bytes = inputs["audio_bytes"]
+            elif "audio" in inputs:
+                audio_bytes = inputs["audio"]
+            else:
+                raise ValueError("No audio found in inputs")
+        else:
+            audio_bytes = inputs
+
+        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
+            tmp.write(audio_bytes)
+            tmp_path = tmp.name
+
+        return self.run_pipeline(tmp_path)

requirements.txt

@@ -0,0 +1,7 @@
+torch
+torchaudio
+pyannote.audio
+transformers
+noisereduce
+jiwer
+librosa

setup.py

@@ -0,0 +1,16 @@
+from setuptools import setup, find_packages
+
+setup(
+    name="asr_diarization",
+    version="0.1.0",
+    packages=find_packages(),
+    install_requires=[
+        "torch",
+        "torchaudio",
+        "pyannote.audio",
+        "transformers",
+        "noisereduce",
+        "jiwer",
+        "librosa"
+    ],
+)