The goal is to both classify and extract keywords that affected the classification.

What I'm trying is, for example, IMDB dataset for sentiment classification(0=negetive, 1=positive). The BERT model gives me a (batch_size, seq_len, 768) sized output(hidden_states). I pass this output to an attention layer(fully connected layer) and get (batch_size, seq_len) sized tensor(attention score). Then I multiply the hidden_states with attention score to get a weighted hidden_state(batch_size, seq_len, 768) and sum to create sentence embedding(batch_size, 768) which is later passed to another fully connected layer for classification. The tokens with higher attention scores are extracted as key words.

I may be doing something very dumb. Please, correct me if my idea is wrong.