SIDetector.py

from transformers import ( # pylint: disable=import-error
    AutoTokenizer,
    AutoModelForSequenceClassification,
    AutoModelForCausalLM,
    pipeline
)

import logging

class SIDetector(object):
    def __init__(self):
        self.classifier = pipeline("sentiment-analysis", model="sentinet/suicidality")

    def forward(self, text: str):
        output = self.classifier(text)[0]
        suicidal = True if output['label'] == 'LABEL_1' else False
        confidence = output['score']
        logging.info(f"Suicidal: {suicidal}, Confidence: {confidence}")
        return suicidal, confidence