from base_model import TextClassifier
import torch
from transformers import pipeline

class PretrainedSentimentAnalyzer(TextClassifier):

    def __init__(self, train_features, train_targets, test_features, test_targets, min_threshold=0.7):

        super().__init__(train_features, train_targets, test_features, test_targets)

    

        device = "cuda" if torch.cuda.is_available() else "cpu"

        self.model = pipeline("text-classification",
                              model="cardiffnlp/twitter-roberta-base-sentiment-latest",
                              device=device)

        self.prediction_map = {'positive' : 'positive',
                               'negative' : 'negative',
                               'neutral' : 'neutral'}

        self.threshold = min_threshold

    def train(self):
        pass

    def predict(self, text_samples:list, inverse_transform:bool, proba:bool=True) -> list:

        predictions = self.model(text_samples, batch_size=128)
        if proba:
            return predictions

        predictions = [self.prediction_map[prediction['label']] if prediction['score'] > self.threshold else 'neutral'
                       for prediction in predictions]

        return predictions