Upload handler.py

handler.py

@@ -0,0 +1,389 @@
+from typing import Dict, List, Any
+from scipy.special import softmax
+import numpy as np
+import weakref
+from utils import (
+    clean_str,
+    clean_str_nopunct,
+    MultiHeadModel,
+    BertInputBuilder,
+    get_num_words,
+    preprocess_transcript_for_eliciting,
+    preprocess_raw_files,
+    post_processing_output_json,
+    compute_student_engagement,
+    compute_talk_time,
+    gpt4_filtering_selection
+)
+import torch
+from transformers import BertTokenizer, BertForSequenceClassification, AutoModelForSequenceClassification, AutoTokenizer
+
+UPTAKE_MODEL='ddemszky/uptake-model'
+QUESTION_MODEL ='ddemszky/question-detection'
+ELICITING_MODEL = 'YaHi/teacher_electra_small'
+
+class UptakeUtterance:
+    def __init__(self, speaker, text, uid=None,
+                 transcript=None, starttime=None, endtime=None, **kwargs):
+        self.speaker = speaker
+        self.text = text
+        self.prev_utt = None
+        self.uid = uid
+        self.starttime = starttime
+        self.endtime = endtime
+        self.transcript = weakref.ref(transcript) if transcript else None
+        self.props = kwargs
+
+        self.uptake = None
+        self.question = None
+
+    def get_clean_text(self, remove_punct=False):
+        if remove_punct:
+            return clean_str_nopunct(self.text)
+        return clean_str(self.text)
+
+    def get_num_words(self):
+        if self.text is None:
+            return 0
+        return get_num_words(self.text)
+
+    def to_dict(self):
+        return {
+            'speaker': self.speaker,
+            'text': self.text,
+            'prev_utt': self.prev_utt,
+            'uid': self.uid,
+            'starttime': self.starttime,
+            'endtime': self.endtime,
+            'uptake': self.uptake,
+            'question':  self.question,
+            **self.props
+        }
+
+    def __repr__(self):
+        return f"Utterance(speaker='{self.speaker}'," \
+               f"text='{self.text}', prev_utt='{self.prev_utt}', uid={self.uid}," \
+               f"starttime={self.starttime}, endtime={self.endtime}, props={self.props})"
+
+class UptakeTranscript:
+    def __init__(self, **kwargs):
+        self.utterances = []
+        self.params = kwargs
+
+    def add_utterance(self, utterance):
+        utterance.transcript = weakref.ref(self)
+        self.utterances.append(utterance)
+
+    def get_idx(self, idx):
+        if idx >= len(self.utterances):
+            return None
+        return self.utterances[idx]
+
+    def get_uid(self, uid):
+        for utt in self.utterances:
+            if utt.uid == uid:
+                return utt
+        return None
+
+    def length(self):
+        return len(self.utterances)
+
+    def to_dict(self):
+        return {
+            'utterances': [utterance.to_dict() for utterance in self.utterances],
+            **self.params
+        }
+
+    def __repr__(self):
+        return f"Transcript(utterances={self.utterances}, custom_params={self.params})"
+    
+class ElicitingUtterance:
+    def __init__(self, speaker, text, starttime, endtime, uid=None, transcript=None, prev_utt=None):
+        self.speaker = speaker
+        self.text = clean_str_nopunct(text)
+        self.uid = uid
+        self.transcript = transcript if transcript else None
+        self.prev_utt = prev_utt
+        self.eliciting = None
+        self.question = None
+        self.starttime = starttime
+        self.endtime = endtime
+
+    def __setitem__(self, key, value):
+        self.__dict__[key] = value
+
+    def get_clean_text(self, remove_punct=False):
+        if remove_punct:
+            return clean_str_nopunct(self.text)
+        return clean_str(self.text)
+
+    def to_dict(self):
+        return {
+            'speaker': self.speaker,
+            'text': self.text,
+            'uid': self.uid,
+            'prev_utt': self.prev_utt,
+            'eliciting': self.eliciting,
+            'question': self.question,
+            'starttime': self.starttime,
+            'endtime': self.endtime,
+        }
+
+
+    def __repr__(self):
+        return f"Utterance(speaker='{self.speaker}'," \
+               f"text='{self.text}', uid={self.uid}, prev_utt={self.prev_utt}, elicting={self.eliciting}, question={self.question}), starttime={self.starttime}, endtime={self.endtime})"
+    
+class ElicitingTranscript:
+    def __init__(self, utterances: List[ElicitingUtterance], tokenizer=None):
+        self.tokenizer = tokenizer
+        self.utterances = []
+        prev_utt = ""
+        prev_utt_teacher = ""
+        prev_speaker = None
+        for utterance in utterances:
+            try:
+                if 'student' in utterance["speaker"]:
+                    utterance["speaker"] = 'student'
+            except:
+                continue
+            if (prev_speaker == 'tutor') and (utterance["speaker"] == 'student'):
+                utterance = ElicitingUtterance(**utterance, transcript=self, prev_utt=prev_utt.text)
+            elif (prev_speaker == 'student') and (utterance["speaker"] == 'tutor'):
+                utterance = ElicitingUtterance(**utterance, transcript=self, prev_utt=prev_utt.text)
+                prev_utt_teacher = utterance.text
+            elif (prev_speaker == 'student') and (utterance["speaker"] == 'student'):
+                try:
+                    utterance = ElicitingUtterance(**utterance, transcript=self, prev_utt=prev_utt_teacher)
+                except:
+                    print("Error on line 159 of handler.py")
+                    print(utterance)
+                    # breakpoint()
+            else:
+                utterance = ElicitingUtterance(**utterance, transcript=self, prev_utt="")
+                if utterance.speaker == 'tutor':
+                    prev_utt_teacher = utterance.text
+            prev_utt = utterance
+            prev_speaker = utterance.speaker
+            self.utterances.append(utterance)
+
+    def __len__(self):
+        return len(self.utterances)
+
+    def __getitem__(self, index):
+        output = self.tokenizer([(self.utterances[index].prev_utt, self.utterances[index].text)], truncation=True)
+        output["speaker"] = self.utterances[index].speaker
+        output["uid"] = self.utterances[index].uid
+        output["prev_utt"] = self.utterances[index].prev_utt
+        output["text"] = self.utterances[index].text
+        return output
+    
+    def to_dict(self):
+        return {
+            'utterances': [utterance.to_dict() for utterance in self.utterances]
+        }
+
+class QuestionModel:
+    def __init__(self, device, tokenizer, input_builder, max_length=300, path=QUESTION_MODEL):
+        print("Loading models...")
+        self.device = device
+        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.to(self.device)
+
+
+    def run_inference(self, transcript):
+        self.model.eval()
+        with torch.no_grad():
+            for i, utt in enumerate(transcript.utterances):
+                if utt.text is None:
+                    utt.question = None
+                    continue
+                if "?" in utt.text:
+                    utt.question = 1
+                else:
+                    text = utt.get_clean_text(remove_punct=True)
+                    instance = self.input_builder.build_inputs([], text,
+                                                               max_length=self.max_length,
+                                                               input_str=True)
+                    output = self.get_prediction(instance)
+                    utt.question = softmax(output["is_question_logits"][0].tolist())[1]
+
+    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"].to(self.device),
+                            attention_mask=instance["attention_mask"].to(self.device),
+                            token_type_ids=instance["token_type_ids"].to(self.device),
+                            return_pooler_output=False)
+        return output
+
+class UptakeModel:
+    def __init__(self, device, tokenizer, input_builder, max_length=120, path=UPTAKE_MODEL):
+        print("Loading models...")
+        self.device = device
+        self.tokenizer = tokenizer
+        self.input_builder = input_builder
+        self.max_length = max_length
+        self.model = MultiHeadModel.from_pretrained(path, head2size={"nsp": 2})
+        self.model.to(self.device)
+
+    def run_inference(self, transcript, min_prev_words, uptake_speaker=None):
+        self.model.eval()
+        prev_num_words = 0
+        prev_utt = None
+        with torch.no_grad():
+            for i, utt in enumerate(transcript.utterances):
+                if ((uptake_speaker is None) or (utt.speaker == uptake_speaker)) and (prev_num_words >= min_prev_words):
+                    textA = prev_utt.get_clean_text(remove_punct=False)
+                    textB = utt.get_clean_text(remove_punct=False)
+                    instance = self.input_builder.build_inputs([textA], textB,
+                                                               max_length=self.max_length,
+                                                               input_str=True)
+                    output = self.get_prediction(instance)
+
+                    utt.uptake = softmax(output["nsp_logits"][0].tolist())[1]
+                    utt.prev_utt = prev_utt.text
+                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].to(self.device)
+
+        output = self.model(input_ids=instance["input_ids"].to(self.device),
+                            attention_mask=instance["attention_mask"].to(self.device),
+                            token_type_ids=instance["token_type_ids"].to(self.device),
+                            return_pooler_output=False)
+        return output
+    
+class ElicitingModel:
+    def __init__(self, device, tokenizer, path=ELICITING_MODEL):
+        print("Loading teacher models...")
+        self.device = device
+        self.tokenizer = tokenizer
+        self.model = AutoModelForSequenceClassification.from_pretrained(path).to(self.device)
+
+    def run_inference(self, dataset):
+        current_batch = 0
+        batch_size = 64
+
+        def generator():
+            while current_batch < len(dataset):
+                yield
+        
+        for _ in generator():
+            # check if the remaining samples are less than the batch size
+            if len(dataset) - current_batch < batch_size:
+                batch_size = len(dataset) - current_batch
+
+            to_pad = [{"input_ids": example["input_ids"][0], "attention_mask": example["attention_mask"][0]} for example in dataset]
+            to_pad = to_pad[current_batch:current_batch + batch_size]
+            batch = self.tokenizer.pad(
+                to_pad,
+                padding=True,
+                max_length=None,
+                pad_to_multiple_of=None,
+                return_tensors="pt",
+            )
+            inputs = batch["input_ids"].to(self.device)
+            attention_mask = batch["attention_mask"].to(self.device)
+            with torch.no_grad():
+                outputs = self.model(inputs, attention_mask=attention_mask)
+            predictions = outputs.logits.argmax(dim=-1).cpu().numpy()
+        
+            for i, prediction in enumerate(predictions):
+                if dataset.utterances[current_batch + i].speaker == 'tutor':
+                    dataset.utterances[current_batch + i]["eliciting"] = prediction
+            current_batch += batch_size
+    
+
+class EndpointHandler():
+    def __init__(self, path="."):
+        print("Loading models...")
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        
+        self.tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
+        self.input_builder = BertInputBuilder(tokenizer=self.tokenizer)
+        self.uptake_model = UptakeModel(self.device, self.tokenizer, self.input_builder)
+        self.question_model = QuestionModel(self.device, self.tokenizer, self.input_builder)
+        self.eliciting_tokenizer = AutoTokenizer.from_pretrained(ELICITING_MODEL)
+        self.eliciting_model = ElicitingModel(self.device, self.tokenizer, path=ELICITING_MODEL)
+
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """
+       data args:
+            inputs (:obj: `list`):
+            List of dicts, where each dict represents an utterance; each utterance object must have a `speaker`,
+            `text` and `uid`and can include list of custom properties
+            parameters (:obj: `dict`)
+       Return:
+            A :obj:`list` | `dict`: will be serialized and returned
+        """
+
+        # get inputs
+        utterances = data.pop("inputs", data)
+        params = data.pop("parameters", None) #TODO: make sure that it includes everything required
+
+        print(params["session_uuid"])
+
+        # pre-processing
+        utterances = preprocess_raw_files(utterances, params)
+
+        # compute student engagement and talk time metrics
+        num_students_engaged, num_students_engaged_talk_only = compute_student_engagement(utterances)
+        tutor_talk_time = compute_talk_time(utterances)
+
+        #TODO: make sure there is some routing going on here based on what session we are at
+        if params["session_type"] == "eliciting":
+            # pre-processing for eliciting
+            utterances_elicting = preprocess_transcript_for_eliciting(utterances)
+            eliciting_transcript = ElicitingTranscript(utterances_elicting, tokenizer=self.tokenizer)
+            self.eliciting_model.run_inference(eliciting_transcript)
+            
+            # Question
+            self.question_model.run_inference(eliciting_transcript)
+
+            transcript_output = eliciting_transcript
+        else:
+            uptake_transcript = UptakeTranscript(filename=params.pop("filename", None))
+            for utt in utterances:
+                uptake_transcript.add_utterance(UptakeUtterance(**utt))
+
+            # Uptake
+            self.uptake_model.run_inference(uptake_transcript, min_prev_words=params['uptake_min_num_words'],
+                                    uptake_speaker=params.pop("uptake_speaker", None))
+
+            # Question
+            self.question_model.run_inference(uptake_transcript)
+            transcript_output = uptake_transcript
+
+        # post-processing
+        model_outputs = post_processing_output_json(transcript_output.to_dict(), params["session_uuid"], params["session_type"])
+        
+        final_output = {}
+        final_output["metrics"] = {"num_students_engaged": num_students_engaged, 
+                                   "num_students_engaged_talk_only": num_students_engaged_talk_only, 
+                                   "tutor_talk_time": tutor_talk_time}
+        
+        if len(model_outputs) > 0:
+            model_outputs = gpt4_filtering_selection(model_outputs, params["session_type"], params["focus_concept"])
+        
+        final_output["model_outputs"] = model_outputs 
+        final_output["event_id"] = params["event_id"]
+
+        import requests
+        webhooks_url = 'https://schoolhouse.world/api/webhooks/stanford-ai-feedback-highlights'
+        response = requests.post(webhooks_url, json=final_output)
+
+        print("Post request sent, here is the response: ", response) 
+
+
+        return final_output