File size: 2,050 Bytes
b964596
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification
)

import torch


MODEL_NAME = "jvomiranda/BERTimbau-Sent-Analysis"


class BertimbauModel:

    def __init__(self):

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

        self.tokenizer = AutoTokenizer.from_pretrained(
            MODEL_NAME
        )

        self.model = (
            AutoModelForSequenceClassification
            .from_pretrained(MODEL_NAME)
            .to(self.device)
        )

        self.model.eval()

    @torch.no_grad()
    def predict(self, text: str):

        encoding = self.tokenizer(
            text,
            return_tensors="pt",
            truncation=True,
            max_length=128
        )

        encoding = {
            key: value.to(self.device)
            for key, value in encoding.items()
        }

        outputs = self.model(**encoding, output_hidden_states=True)

        probabilities = torch.softmax(
            outputs.logits,
            dim=1
        )[0]

        prediction = torch.argmax(
            probabilities
        ).item()

        tokens = self.tokenizer.convert_ids_to_tokens(
            encoding["input_ids"][0]
        )

        token_ids = (
            encoding["input_ids"][0]
            .cpu()
            .tolist()
        )

        embeddings = outputs.hidden_states[-1][0]

        representation = []

        for token, vector in zip(tokens, embeddings):

            representation.append({
                "token": token,
                "vector": [
                    round(float(x), 2)
                    for x in vector[:5]
                ]
            })

        return {

            "prediction": prediction,

            "probabilities": {

                0: float(probabilities[0]),

                1: float(probabilities[1])

            },

            "tokens": tokens,

            "token_ids": token_ids,

            "representation": representation

        }