batch prediction, adding post processing

tim-roethig-db · May 25, 2024 · b2dcf26 · b2dcf26
1 parent dde7f80
commit b2dcf26
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 1 deletion.
diff --git a/amondin/__init__.py b/amondin/__init__.py
@@ -5,3 +5,4 @@
 from .speech2text import speech2text
 from .tools import get_secret
 from .main import transcribe
+from .post_processing import merge_rows_consecutive_speaker
diff --git a/amondin/main.py b/amondin/main.py
@@ -8,6 +8,7 @@
 from amondin.tools import get_secret
 from amondin.segment_speakers import segment_speakers
 from amondin.speech2text import speech2text
+from amondin.post_processing import merge_rows_consecutive_speaker
 
 
 def transcribe(
@@ -70,6 +71,8 @@ def transcribe(
 
     transcript = pd.DataFrame(segments)
 
+    transcript = merge_rows_consecutive_speaker(transcript)
+
     # save transcript
     print(transcript.to_markdown(index=False))
     if output_file_path.endswith(".csv"):

diff --git a/amondin/post_processing.py b/amondin/post_processing.py
@@ -7,7 +7,8 @@ def merge_rows_consecutive_speaker(transcript: pd.DataFrame) -> pd.DataFrame:
     print(transcript.to_markdown())
 
     transcript = transcript.groupby(['speaker_group', 'speaker']).agg({
-        'time_stamp': lambda x: ' '.join(x),
+        'start': "min",
+        "end": "max",
         'text': lambda x: ' '.join(x)
     }).reset_index()