From 4cb774007d62a24c5f8b85ade5e3d1dbd7ce4c09 Mon Sep 17 00:00:00 2001
From: Jaikinator <schmieder.jacob@web.de>
Date: Mon, 19 Jun 2023 15:06:16 +0200
Subject: [PATCH] added save function to export to json

---
 autotranscript/diarisation.py | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/autotranscript/diarisation.py b/autotranscript/diarisation.py
index 5b71f88..070fc2d 100644
--- a/autotranscript/diarisation.py
+++ b/autotranscript/diarisation.py
@@ -1,8 +1,9 @@
-from pyannote.audio import Pipeline
+from .audio import Pipeline
 from pyannote.audio.pipelines.speaker_diarization import SpeakerDiarization
 from torch import Tensor
 import os
 from typing import TypeVar, Union
+import json
 
 Annotation = TypeVar('Annotation') 
 
@@ -88,6 +89,18 @@ class Diariser:
             diarization_output["speakers"].append(outp[2])
         return diarization_output
     
+    def save(self, path : str, *args, **kwargs) -> None:
+        """
+        Save diarization output to a file
+
+        :param path: path to save file
+        :type path: str
+        """
+        with open(path, "w") as f:
+            json.dump(self.transcript, f, *args, **kwargs)
+        
+        
+    
     @staticmethod
     def _get_token():
         # check ig .pyannotetoken.txt exists