Upload folder using huggingface_hub

asr_diarization/pipeline.py

@@ -48,47 +48,9 @@ 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)
-        all_word_segments = []
+        merged_segments = []
         speaker_segments = {}
 
         for seg in diar_json:
@@ -99,45 +61,32 @@ 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"]:
-                    # 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())
+                    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)
+            print("Sample merged segment:", merged_segments[0])
+
+
+            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):