add formatting

1 files changed, 17 insertions, 3 deletions

diff --git a/asr.py b/asr.py
index e598e4b..54b83eb 100644
--- a/asr.py
+++ b/asr.py
@@ -8,10 +8,19 @@ from tqdm import tqdm
 
 # converts TimestampedResult to JSON-serializable dictionary
 def result_to_dict(result, offset_seconds):
+    token_map = [
+        [token, round(ts + offset_seconds, 2)]
+        for token, ts in zip(result.tokens, result.timestamps)
+    ]
+
+    start_time = token_map[0][1] if token_map else offset_seconds
+    end_time = token_map[-1][1] if token_map else offset_seconds
+
     return {
         "text": result.text,
-        "tokens": result.tokens,
-        "timestamps": [ts + offset_seconds for ts in result.timestamps]
+        "start_time": start_time,
+        "end_time": end_time,
+        "token_map": token_map
     }
 
 def main(input_file, output_file):
@@ -56,7 +65,12 @@ def main(input_file, output_file):
             transcript_results = model.recognize(np.copy(chunk))
             
             if transcript_results:
-                for result in transcript_results:
+                if not isinstance(transcript_results, list):
+                    results_list = [transcript_results]
+                else:
+                    results_list = transcript_results
+
+                for result in results_list:
                     full_transcript_data.append(result_to_dict(result, chunk_offset_seconds))
         except Exception as e:
             print(f"  Error processing chunk {i + 1}: {e}")