tweet_pipeline.py

import os
from typing import Dict

from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    TextClassificationPipeline,
)


class TweetPipeline:
    def __init__(self):
        self.emotion_pipeline = TextClassificationPipeline(
            model=AutoModelForSequenceClassification.from_pretrained(
                "elozano/tweet_emotion_eval"
            ),
            tokenizer=AutoTokenizer.from_pretrained("elozano/tweet_emotion_eval"),
        )
        self.offensive_pipeline = TextClassificationPipeline(
            model=AutoModelForSequenceClassification.from_pretrained(
                "elozano/tweet_offensive_eval"
            ),
            tokenizer=AutoTokenizer.from_pretrained("elozano/tweet_offensive_eval"),
        )
        self.sentiment_pipeline = TextClassificationPipeline(
            model=AutoModelForSequenceClassification.from_pretrained(
                "elozano/tweet_sentiment_eval"
            ),
            tokenizer=AutoTokenizer.from_pretrained("elozano/tweet_sentiment_eval"),
        )

    def __call__(self, text: str) -> Dict[str, str]:
        return {
            "emotion": self.emotion_pipeline(text)[0]["label"],
            "offensive": self.offensive_pipeline(text)[0]["label"],
            "sentiment": self.sentiment_pipeline(text)[0]["label"],
        }