Spaces:
Runtime error
Runtime error
| 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"], | |
| } | |