Upload folder using huggingface_hub

asr_diarization/pipeline.py

@@ -48,9 +48,47 @@ class ASR_Diarization:
             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:
+    #         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)
+    #         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({
+    #                     "tag": "w",
+    #                     "start": start_ts,
+    #                     "end": end_ts,
+    #                     "text": word_info["text"]
+    #                 })
+
+    #         seg_dict = {
+    #             "speaker": spk,
+    #             "segment_start": segment_start,
+    #             "segment_end": segment_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_transcription(self, audio_path, diar_json):
         audio, sr = torchaudio.load(audio_path)
-        merged_segments = []
+        all_word_segments = []
         speaker_segments = {}
 
         for seg in diar_json:
@@ -61,30 +99,45 @@ class ASR_Diarization:
             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({
-                        "tag": "w",
-                        "start": start_ts,
-                        "end": end_ts,
-                        "text": word_info["text"]
-                    })
-
-            seg_dict = {
-                "speaker": spk,
-                "segment_start": segment_start,
-                "segment_end": segment_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())
+                    # Each word or token gets its own mini segment
+                    start_ts, end_ts = None, None
+
+                    if isinstance(word_info.get("timestamp"), (list, tuple)):
+                        start_ts, end_ts = word_info["timestamp"]
+                    elif isinstance(word_info.get("timestamp"), (float, int)):
+                        start_ts = word_info["timestamp"]
+                        end_ts = start_ts
+
+                    if start_ts is None:
+                        continue
+
+                    # Shift timestamps to align with full audio
+                    abs_start = segment_start + start_ts
+                    abs_end = segment_start + end_ts
+
+                    word_segment = {
+                        "speaker": spk,
+                        "segment_start": abs_start,
+                        "segment_end": abs_end,
+                        "tokens": [
+                            {
+                                "tag": "w",
+                                "start": abs_start,
+                                "end": abs_end,
+                                "text": word_info["text"].strip()
+                            }
+                        ]
+                    }
+
+                    all_word_segments.append(word_segment)
+
+                    if spk not in speaker_segments:
+                        speaker_segments[spk] = []
+                    speaker_segments[spk].append(word_segment)
+
+        return all_word_segments, list(speaker_segments.keys())
 
     def run_pipeline(self, audio_path, output_dir=None, base_name=None,
                      ref_rttm=None, ref_json=None):