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