Add json converting

handler.py

@@ -12,6 +12,7 @@ import webvtt
 from datetime import datetime
 import torch
 import spacy
+import json
 
 
 nlp = spacy.load("en_core_web_sm")
@@ -199,6 +200,18 @@ class EndpointHandler():
         print(utterances_list)
         return utterances_list
         
+
+    def utterance_list_to_json(self, utterances: List[Utterance], use_prior_text: bool) -> List[str]:
+        formatted = []
+        for utterance in utterances:
+            if not use_prior_text:
+                formatted.append({'speaker': utterance.speaker, 'data': utterance.text, 'time': utterance.starttime})
+            else:
+                formatted.append([{'speaker': utterance.prev_prev_utterance.speaker, 'data': utterance.prev_prev_utterance.text, 'time': utterance.prev_prev_utterance.starttime},
+                                  {'speaker': utterance.prev_utterance.speaker, 'data': utterance.prev_utterance.text, 'time': utterance.prev_utterance.starttime},
+                                  {'speaker': utterance.speaker, 'data': utterance.text, 'time': utterance.starttime}])
+        return formatted
+
     def __call__(self, data: str) -> List[Dict[str, Any]]:
         ''' data_file is a str pointing to filename of type .vtt '''
 
@@ -208,9 +221,12 @@ class EndpointHandler():
 
         if data_file is None:
             raise ValueError("no data file provided")
-        
+
+        full_transcript = self.process_vtt_transcript(data_file)
         utterances_list = []
-        for utterance in self.process_vtt_transcript(data_file):
+        utterances_indexes = [] # entry corresponds to utterance in full_transcript
+        for i in range(len(full_transcript)):
+            utterance = full_transcript[i]
             #TODO: filter out to only have SL utterances
             if model_id == 'eliciting':
                 utterance_str, is_list = self.eliciting_utterance_to_str(utterance)
@@ -223,11 +239,16 @@ class EndpointHandler():
             
             if is_list == 'list':
                 utterances_list.extend(utterance_str)
+                for j in range(len(utterance_str)):
+                    utterances_indexes.append(i)
             else:
                 utterances_list.append(utterance_str)
+                utterances_indexes.append(i)
 
+        talk_move = ""
         cuda_available = torch.cuda.is_available()
         if model_id == 'eliciting':
+            talk_move = 'getIdeas'
             self.model = ClassificationModel(
                 "roberta", "aekupor/eliciting", use_cuda=cuda_available
             )
@@ -240,6 +261,8 @@ class EndpointHandler():
                 "roberta", "aekupor/probing", use_cuda=cuda_available
             )
         elif model_id == 'adding_on':
+            # TODO: combine adding on and others
+            talk_move = 'buildIdeas'
             self.model = ClassificationModel(
                 "roberta", "aekupor/adding_on", use_cuda=cuda_available
             )
@@ -247,6 +270,31 @@ class EndpointHandler():
             raise ValueError(f"model_id: {model_id} is not valid. Available models are: {list(self.multi_model.keys())}")
         
         predictions, _ = self.model.predict(utterances_list)
+
+        # json formating
+        full_transcript_json = json.dumps(self.utterance_list_to_json(full_transcript, False), separators=(',', ':'))
+        print("FULL TRANSCRIPT")
+        print(full_transcript_json)
+
+        utterance_talk_moves = set()
+        for i in range(len(predictions)):
+            if predictions[i] == 1:
+                utterance_talk_moves.add(full_transcript[utterances_indexes[i]])
+
+        utterance_talk_moves_json = ''
+        if model_id == 'elicting' or model_id == 'connecting':
+            utterance_talk_moves_json = json.dumps(self.utterance_list_to_json(utterance_talk_moves, False), separators=(',', ':'))
+        elif model_id == 'adding_on':
+            utterance_talk_moves_json = json.dumps(self.utterance_list_to_json(utterance_talk_moves, True), separators=(',', ':'))
+            
+        print("TALK MOVES FOUND")
+        print(utterance_talk_moves_json)
+
+        print("TALK MOVE")
+        print(talk_move)
+
+        print("NUM TALK MOVES")
+        print(len(utterance_talk_moves))
         
         return predictions