File size: 768 Bytes
3fe0726
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
# Load model directly
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from news_scraper.nlp_models.basemodel import BaseModel


class DistilRoBERTa(BaseModel):
    """
    A class for performing sentiment analysis using the FinBERT model.
    
    https://huggingface.co/mrm8488/distilroberta-finetuned-financial-news-sentiment-analysis
    """
    
    def __init__(self):
        super().__init__("DistilRoBERTa")
        self.tokenizer = AutoTokenizer.from_pretrained("mrm8488/distilroberta-finetuned-financial-news-sentiment-analysis")
        self.model = AutoModelForSequenceClassification.from_pretrained("mrm8488/distilroberta-finetuned-financial-news-sentiment-analysis")
        self.model.to(self.device)
        self.model.eval()