output data in jsonl

1 files changed, 6 insertions, 4 deletions

diff --git a/main.py b/main.py
index af76aca..283ea9a 100644
--- a/main.py
+++ b/main.py
@@ -3,6 +3,7 @@ import re
 import sys
 import backoff
 import openai
+import json
 from datetime import datetime, timedelta
 from tqdm import tqdm
 
@@ -263,21 +264,22 @@ def process_file(input_file: str, output_file: str):
         if ad_timestamps_str:
             # Split by newlines in case the model returns multiple segments
             found_timestamps = [ts.strip() for ts in ad_timestamps_str.split('\n') if ts.strip()]
-            all_ad_timestamps.extend(found_timestamps)
+            all_ad_timestamps.append({"id": f"{"".join(input_file.split(".")[:-1])}_chunk{i}", f"text": chunk_text_for_api, "target": "\n".join(found_timestamps)})
 
     with open(output_file, "w") as f:
-        f.write("\n".join(all_ad_timestamps))
+        for s in all_ad_timestamps:
+            f.write(json.dumps(s) + "\n")
 
 def main():
     paths = sys.argv[1:]
     for p in paths:
-        out = "".join(p.split(".")[:-1]) + ".gpt"
+        out = "".join(p.split(".")[:-1]) + ".jsonl"
         if os.path.isfile(out):
             paths.remove(p)
     print(paths)
 
     for p in tqdm(paths):
-        out = "".join(p.split(".")[:-1]) + ".gpt"
+        out = "".join(p.split(".")[:-1]) + ".jsonl"
         process_file(p, out)
 
 if __name__ == "__main__":