文書分類における古典的手法とBERTの判断根拠の比較

$ wget https://archive.ics.uci.edu/ml/machine-learning-databases/00228/smsspamcollection.zip
$ unzip smsspamcollection.zip

from sklearn.model_selection import train_test_split
import pandas as pd

SEED = 0

spam_df = pd.read_csv("SMSSpamCollection", sep='\t', header=None)  # データの読み込み

# ラベルと文章を分ける
labels = spam_df[0].values
sentences = spam_df[1].values

label_dic = {'ham': 0, 'spam': 1}  # spamを真値とする
label_dic_inv = {v: k for k, v in label_dic.items()}
label_ids = [label_dic[i] for i in labels]

# 7:3に学習データとテストデータを分割する
train_sentence, test_sentence, y_train, y_test = train_test_split(sentences, label_ids, test_size=0.3, random_state=SEED, stratify=label_ids)

from sklearn.feature_extraction.text import CountVectorizer

vectorizer = CountVectorizer()
X_train = vectorizer.fit_transform(train_sentence)
X_test = vectorizer.transform(test_sentence)

from sklearn.linear_model import LogisticRegression
from sklearn.metrics import confusion_matrix

lr = LogisticRegression(random_state=SEED, n_jobs=-1)
lr.fit(X_train, y_train)

y_pred = lr.predict(X_test)
cm = confusion_matrix(y_test, y_pred)

pd.DataFrame(cm,columns=['Predicted ham', 'Predicted spam'], index=['Actual ham', 'Actual spam'])

from IPython.display import display, HTML

# 赤くハイライトする
def highlight_r(word, attn):
  html_color = '#%02X%02X%02X' % (255, int(255*(1 - attn)), int(255*(1 - attn)))
  return '<span style="background-color: {}">{}</span>'.format(html_color, word)

# 青くハイライトする
def highlight_b(word, attn):
  html_color = '#%02X%02X%02X' % (int(255*(1 - attn)), int(255*(1 - attn)), 255)
  return '<span style="background-color: {}">{}</span>'.format(html_color, word)

def show_lr_explaination(check_idx):
    # 単語と説明変数の値の辞書
    coef_dic = {j: i for i, j in zip(lr.coef_[0], vectorizer.get_feature_names())}
        
    # 対象の文章の単語の説明変数の値を確認していく
    texts = test_sentence[check_idx].split()
    scores = []
    for w in texts:
        try:
            s = coef_dic[w]
        except KeyError:
            s = 0  # 対象外の単語は0を割り当てる
        scores.append(s)

    # 文章をハイライトしていく
    html_outputs = []
    for word, attn in zip(texts, scores):
        if attn < 0:
            html_outputs.append(highlight_b(word, attn*-1))
        else:
            html_outputs.append(highlight_r(word, attn))

    # 結果を表示
    display(HTML(' '.join(html_outputs)))

class SpamBert(nn.Module):
    def __init__(self, model_type, tokenizer):
        super(SpamBert, self).__init__()

        bert_conf = BertConfig(model_type, output_hidden_states=False, output_attentions=True)
        bert_conf.vocab_size = tokenizer.vocab_size

        self.bert = AutoModel.from_pretrained(model_type, config=bert_conf)
        self.fc = nn.Linear(bert_conf.hidden_size, 1)

    def forward(self, ids, mask, token_type_ids):
        out = self.bert(ids, attention_mask=mask, token_type_ids=token_type_ids)
        h = out['pooler_output']
        a = out['attentions']
        h = nn.ReLU()(h)
        h = self.fc(h)
        h = h[:, 0]
        a = a[-1].sum(1)[:, 0, :]
        return h, a

def show_bert_explaination(check_idx):
    for idx, d in enumerate(test_dataset):
        if idx == check_idx:
            break

    input_ids = d["input_ids"].to(device).unsqueeze(0)
    attention_mask = d["attention_mask"].to(device).unsqueeze(0)
    token_type_ids = d["token_type_ids"].to(device).unsqueeze(0)
    target = d["target"].to(device)

    with torch.no_grad():
        output, attention = model(input_ids, attention_mask, token_type_ids)

    attention = attention.cpu()[0].numpy()
    attention_mask = attention_mask.cpu()[0].numpy()
    attention = attention[attention_mask == 1][1:-1]

    ids = input_ids.cpu()[0][attention_mask == 1][1:-1].tolist()
    tokens = TOKENIZER.convert_ids_to_tokens(ids)

    html_outputs = []
    
    for word, attn in zip(tokens, attention):
        html_outputs.append(highlight_r(word, attn))
    
    display(HTML(' '.join(html_outputs)))

def predictor(texts):
    tok = TOKENIZER.batch_encode_plus(texts, padding=True)
    input_ids = torch.tensor(tok['input_ids']).to(device)
    attention_mask = torch.tensor(tok['attention_mask']).to(device)
    token_type_ids = torch.tensor(tok['token_type_ids']).to(device)

    with torch.no_grad():
        output, _ = model(input_ids, attention_mask, token_type_ids)
    probas = output.sigmoid().cpu().numpy()

    # 出力は[negativeの予測確率, positiveの予測確率]で出す
    return np.vstack([1 - probas, probas]).T

x = np.arange(len(tokens))

scaler = MinMaxScaler()
color_arr = scaler.fit_transform(np.array(cam_w).reshape(-1, 1))
colorlist = [[1.0, 1-min([c, 1.0]), 1.0] for c in color_arr.T[0]]


width = 0.35

fig, ax = plt.subplots(figsize=(24,4))

rect = ax.bar(x, cam_w, width, color=colorlist)
ax.set_xticks(x)
ax.set_xticklabels(tokens, rotation=45)
plt.ylim(0.027, 0.035)

plt.show()