add focusing questions

handler.py

@@ -13,9 +13,11 @@ from transformers import BertTokenizer, BertForSequenceClassification
 
 transformers.logging.set_verbosity_debug()
 
-UPTAKE_MODEL='ddemszky/uptake-model'
-REASONING_MODEL ='ddemszky/student-reasoning'
-QUESTION_MODEL ='ddemszky/question-detection'
+UPTAKE_MODEL = 'ddemszky/uptake-model'
+REASONING_MODEL = 'ddemszky/student-reasoning'
+QUESTION_MODEL = 'ddemszky/question-detection'
+FOCUSING_QUESTION_MODEL = 'ddemszky/focusing-questions'
+
 
 class Utterance:
     def __init__(self, speaker, text, uid=None,
@@ -31,6 +33,7 @@ class Utterance:
         self.uptake = None
         self.reasoning = None
         self.question = None
+        self.focusing_question = None 
 
     def get_clean_text(self, remove_punct=False):
         if remove_punct:
@@ -50,6 +53,7 @@ class Utterance:
             'uptake': self.uptake,
             'reasoning': self.reasoning,
             'question':  self.question,
+            'focusingquestion': self.focusing_question,
             **self.props
         }
 
@@ -58,6 +62,7 @@ class Utterance:
                f"text='{self.text}', uid={self.uid}," \
                f"starttime={self.starttime}, endtime={self.endtime}, props={self.props})"
 
+
 class Transcript:
     def __init__(self, **kwargs):
         self.utterances = []
@@ -90,6 +95,7 @@ class Transcript:
     def __repr__(self):
         return f"Transcript(utterances={self.utterances}, custom_params={self.params})"
 
+
 class QuestionModel:
     def __init__(self, device, tokenizer, input_builder, max_length=300, path=QUESTION_MODEL):
         print("Loading models...")
@@ -97,10 +103,10 @@ class QuestionModel:
         self.tokenizer = tokenizer
         self.input_builder = input_builder
         self.max_length = max_length
-        self.model = MultiHeadModel.from_pretrained(path, head2size={"is_question": 2})
+        self.model = MultiHeadModel.from_pretrained(
+            path, head2size={"is_question": 2})
         self.model.to(self.device)
 
-
     def run_inference(self, transcript):
         self.model.eval()
         with torch.no_grad():
@@ -114,12 +120,14 @@ class QuestionModel:
                                                                input_str=True)
                     output = self.get_prediction(instance)
                     print(output)
-                    utt.question = np.argmax(output["is_question_logits"][0].tolist())
+                    utt.question = np.argmax(
+                        output["is_question_logits"][0].tolist())
 
     def get_prediction(self, instance):
         instance["attention_mask"] = [[1] * len(instance["input_ids"])]
         for key in ["input_ids", "token_type_ids", "attention_mask"]:
-            instance[key] = torch.tensor(instance[key]).unsqueeze(0)  # Batch size = 1
+            instance[key] = torch.tensor(
+                instance[key]).unsqueeze(0)  # Batch size = 1
             instance[key].to(self.device)
 
         output = self.model(input_ids=instance["input_ids"],
@@ -128,6 +136,7 @@ class QuestionModel:
                             return_pooler_output=False)
         return output
 
+
 class ReasoningModel:
     def __init__(self, device, tokenizer, input_builder, max_length=128, path=REASONING_MODEL):
         print("Loading models...")
@@ -152,7 +161,8 @@ class ReasoningModel:
     def get_prediction(self, instance):
         instance["attention_mask"] = [[1] * len(instance["input_ids"])]
         for key in ["input_ids", "token_type_ids", "attention_mask"]:
-            instance[key] = torch.tensor(instance[key]).unsqueeze(0)  # Batch size = 1
+            instance[key] = torch.tensor(
+                instance[key]).unsqueeze(0)  # Batch size = 1
             instance[key].to(self.device)
 
         output = self.model(input_ids=instance["input_ids"],
@@ -160,6 +170,7 @@ class ReasoningModel:
                             token_type_ids=instance["token_type_ids"])
         return output
 
+
 class UptakeModel:
     def __init__(self, device, tokenizer, input_builder, max_length=120, path=UPTAKE_MODEL):
         print("Loading models...")
@@ -184,14 +195,16 @@ class UptakeModel:
                                                                input_str=True)
                     output = self.get_prediction(instance)
 
-                    utt.uptake = int(softmax(output["nsp_logits"][0].tolist())[1] > .8)
+                    utt.uptake = int(
+                        softmax(output["nsp_logits"][0].tolist())[1] > .8)
                 prev_num_words = utt.get_num_words()
                 prev_utt = utt
 
     def get_prediction(self, instance):
         instance["attention_mask"] = [[1] * len(instance["input_ids"])]
         for key in ["input_ids", "token_type_ids", "attention_mask"]:
-            instance[key] = torch.tensor(instance[key]).unsqueeze(0)  # Batch size = 1
+            instance[key] = torch.tensor(
+                instance[key]).unsqueeze(0)  # Batch size = 1
             instance[key].to(self.device)
 
         output = self.model(input_ids=instance["input_ids"],
@@ -201,6 +214,44 @@ class UptakeModel:
         return output
 
 
+
+class FocusingQuestionModel:
+    def __init__(self, device, tokenizer, input_builder, max_length=128, path=FOCUSING_QUESTION_MODEL):
+        print("Loading models...")
+        self.device = device
+        self.tokenizer = tokenizer
+        self.input_builder = input_builder
+        self.model = BertForSequenceClassification.from_pretrained(path)
+        self.model.to(self.device)
+        self.max_length = max_length
+
+    def run_inference(self, transcript, min_focusing_words=0, uptake_speaker=None):
+        self.model.eval()
+        with torch.no_grad():
+            for i, utt in enumerate(transcript.utterances):
+                if utt.speaker != uptake_speaker or uptake_speaker is None:
+                    utt.focusing_question = None
+                    continue
+                if utt.get_num_words() < min_focusing_words:
+                    utt.focusing_question = None
+                    continue
+                instance = self.input_builder.build_inputs([], utt.text, max_length=self.max_length, input_str=True)
+                output = self.get_prediction(instance)
+                utt.focusing_question = np.argmax(output["logits"][0].tolist())
+
+    def get_prediction(self, instance):
+        instance["attention_mask"] = [[1] * len(instance["input_ids"])]
+        for key in ["input_ids", "token_type_ids", "attention_mask"]:
+            instance[key] = torch.tensor(
+                instance[key]).unsqueeze(0)  # Batch size = 1
+            instance[key].to(self.device)
+
+        output = self.model(input_ids=instance["input_ids"],
+                            attention_mask=instance["attention_mask"],
+                            token_type_ids=instance["token_type_ids"])
+        return output        
+
+
 class EndpointHandler():
     def __init__(self, path="."):
         print("Loading models...")
@@ -231,18 +282,26 @@ class EndpointHandler():
             transcript.add_utterance(Utterance(**utt))
 
         print("Running inference on %d examples..." % transcript.length())
-
+        uptake_speaker = params.pop("uptake_speaker", None)
         # Uptake
-        uptake_model = UptakeModel(self.device, self.tokenizer, self.input_builder)
+        uptake_model = UptakeModel(
+            self.device, self.tokenizer, self.input_builder)
         uptake_model.run_inference(transcript, min_prev_words=params['uptake_min_num_words'],
-                                   uptake_speaker=params.pop("uptake_speaker", None))
+                                   uptake_speaker=uptake_speaker)
 
         # Reasoning
-        reasoning_model = ReasoningModel(self.device, self.tokenizer, self.input_builder)
+        reasoning_model = ReasoningModel(
+            self.device, self.tokenizer, self.input_builder)
         reasoning_model.run_inference(transcript)
 
         # Question
-        question_model = QuestionModel(self.device, self.tokenizer, self.input_builder)
+        question_model = QuestionModel(
+            self.device, self.tokenizer, self.input_builder)
         question_model.run_inference(transcript)
 
+        # Focusing Question
+        focusing_question_model = FocusingQuestionModel(
+            self.device, self.tokenizer, self.input_builder)
+        focusing_question_model.run_inference(transcript, uptake_speaker=uptake_speaker)
+
         return transcript.to_dict()