File size: 1,005 Bytes
f453042 c1fe34c f453042 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 | 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
|