| from transformers import pipeline | |
| ON_TASK_MODEL = 'edsi-umd/on-task-bert' | |
| ON_TASK_LABEL = 'on_task' # from the model's config.json id2label: {"0": "off_task", "1": "on_task"} | |
| class OnTaskAnalyser: | |
| def __init__(self, model_path=ON_TASK_MODEL, max_length=256): | |
| self.pipe = pipeline("text-classification", model=model_path, | |
| truncation=True, max_length=max_length) | |
| def predict_one(self, text: str): | |
| return self.pipe(text)[0] # {'label': ..., 'score': ...} | |
| def run_analysis(self, transcript, uptake_speaker=None): | |
| """Mutate transcript utterances by setting on_task for student utterances.""" | |
| for utt in transcript.utterances: | |
| if uptake_speaker is not None and utt.speaker == uptake_speaker: | |
| continue # model trained on student utterances only | |
| if not utt.text or not utt.text.strip(): | |
| continue | |
| utt.on_task = self.predict_one(utt.text) | |
| return transcript | |