Upload folder using huggingface_hub

asr_diarization/pipeline.py

@@ -44,7 +44,7 @@ class ASR_Diarization:
     def run_diarization(self, audio_path):
         diarization = self.diar_pipeline(audio_path)
         return [
-            {"segment_start": t.start, "segment_end": t.end, "speaker": spk}
+            {"start": t.start, "end": t.end, "speaker": spk}
             for t, _, spk in diarization.itertracks(yield_label=True)
         ]
 
@@ -54,8 +54,8 @@ class ASR_Diarization:
         speaker_segments = {}
 
         for seg in diar_json:
-            segment_start, segment_end, spk = seg["segment_start"], seg["segment_end"], seg["speaker"]
-            start_sample, end_sample = int(segment_start * sr), int(segment_end * sr)
+            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)
@@ -66,21 +66,19 @@ class ASR_Diarization:
                 for word_info in result["chunks"]:
                     start_ts, end_ts = word_info.get("timestamp", (None, None)) or (None, None)
                     tokens.append({
-                        "tag": "w",
                         "start": start_ts,
                         "end": end_ts,
-                        "text": word_info["text"]
+                        "text": word_info["text"],
+                        "tag": "w"
                     })
 
             seg_dict = {
                 "speaker": spk,
-                "segment_start": segment_start,
-                "segment_end": segment_end,
+                "start": start,
+                "end": end,
                 "tokens": tokens
             }
             merged_segments.append(seg_dict)
-            print("Sample merged segment:", merged_segments[0])
-
 
             if spk not in speaker_segments:
                 speaker_segments[spk] = []
@@ -101,8 +99,8 @@ class ASR_Diarization:
             with open(rttm_path, "w") as f:
                 for seg in diar_json:
                     f.write(
-                        f"SPEAKER {base_name} 1 {seg['segment_start']:.6f} "
-                        f"{seg['segment_end']-seg['segment_start']:.6f} <NA> <NA> "
+                        f"SPEAKER {base_name} 1 {seg['start']:.6f} "
+                        f"{seg['end']-seg['start']:.6f} <NA> <NA> "
                         f"{seg['speaker']} <NA>\n"
                     )