【🔰Huggingface Transformers入門⑦】文章分類モデルを作成する（2）　〜Trainerクラスとファインチューニング〜

from google.colab import drive 
drive.mount('/content/drive')
!mkdir -p '/content/drive/My Drive/huggingface_transformers_demo/'
%cd '/content/drive/My Drive/huggingface_transformers_demo/'

!pip install transformers fugashi ipadic sentencepiece datasets

!wget https://www.rondhuit.com/download/ldcc-20140209.tar.gz
!tar zxvf ldcc-20140209.tar.gz

import os
import pandas as pd

# タイトルリストの取得
def get_title_list(path):
    title_list = []
    filenames = os.listdir(path)
    for filename in filenames:
        # ファイルの読み込み
        with open(path+filename) as f:
            title = f.readlines()[2].strip()
            title_list.append(title)
    return title_list

# データセットの生成（各ニュース記事のタイトルを取得）
df = pd.DataFrame(columns=['label', 'sentence'])

#0
title_list = get_title_list('text/dokujo-tsushin/')
for title in title_list:
    df = df.append({'label':0 , 'sentence':title}, ignore_index=True)

#1
title_list = get_title_list('text/it-life-hack/')
for title in title_list:
    df = df.append({'label':1 , 'sentence':title}, ignore_index=True)

#2
title_list = get_title_list('text/sports-watch/')
for title in title_list:
   df = df.append({'label':2 , 'sentence':title}, ignore_index=True)

#3
title_list = get_title_list('text/kaden-channel/')
for title in title_list:
   df = df.append({'label':3 , 'sentence':title}, ignore_index=True)

#4
title_list = get_title_list('text/livedoor-homme/')
for title in title_list:
   df = df.append({'label':4 , 'sentence':title}, ignore_index=True)
#5
title_list = get_title_list('text/movie-enter/')
for title in title_list:
   df = df.append({'label':5 , 'sentence':title}, ignore_index=True)

#6
title_list = get_title_list('text/peachy/')
for title in title_list:
   df = df.append({'label':6 , 'sentence':title}, ignore_index=True)

#7
title_list = get_title_list('text/smax/')
for title in title_list:
   df = df.append({'label':7 , 'sentence':title}, ignore_index=True)

#8
title_list = get_title_list('text/topic-news/')
for title in title_list:
   df = df.append({'label':8 , 'sentence':title}, ignore_index=True)

# データをシャッフルする
df = df.sample(frac=1)

# 上記で取得したデータを学習用と検証用に分割してCSV出力
num = len(df)
df[:int(num*0.8)].to_csv('./text/news_train.csv', sep=',', index=False)
df[int(num*0.8):].to_csv('./text/news_val.csv', sep=',', index=False)

from transformers import AutoModelForSequenceClassification, AutoTokenizer
from transformers import TrainingArguments
from transformers import Trainer
from datasets import Dataset
from sklearn.metrics import accuracy_score, f1_score
from sklearn.metrics import ConfusionMatrixDisplay, confusion_matrix
import torch
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

train_df = pd.read_csv("./text/news_train.csv",encoding='UTF-8')
validation_df = pd.read_csv("./text/news_val.csv",encoding='UTF-8')

train_dataset = Dataset.from_pandas(train_df)
validation_dataset = Dataset.from_pandas(validation_df)

from datasets import DatasetDict
dataset = DatasetDict({
    "train": train_dataset,
    "validation": validation_dataset,
})

train_df

# トークナイザの取得
tokenizer = AutoTokenizer.from_pretrained("cl-tohoku/bert-base-japanese-whole-word-masking")

# モデルの取得
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
num_labels = 9
model = (AutoModelForSequenceClassification
    .from_pretrained("cl-tohoku/bert-base-japanese-whole-word-masking", num_labels=num_labels)
    .to(device))

# トークナイザ処理
def tokenize(batch):
    return tokenizer(batch["sentence"], padding=True, truncation=True)
dataset_encoded = dataset.map(tokenize, batched=True, batch_size=None)

# 評価指標の定義
def compute_metrics(pred):
    labels = pred.label_ids
    preds = pred.predictions.argmax(-1)
    f1 = f1_score(labels, preds, average="weighted")
    acc = accuracy_score(labels, preds)
    return {"accuracy": acc, "f1": f1}

# 学習の準備
batch_size = 16
logging_steps = len(dataset_encoded["train"]) // batch_size
model_name = f"classification"
training_args = TrainingArguments(
    output_dir=model_name,
    num_train_epochs=2,
    learning_rate=2e-5,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    weight_decay=0.01,
    evaluation_strategy="epoch",
    disable_tqdm=False,
    logging_steps=logging_steps,
    push_to_hub=False,
    log_level="error",
)

# 学習
trainer = Trainer(
    model=model, args=training_args,
    compute_metrics=compute_metrics,
    train_dataset=dataset_encoded["train"],
    eval_dataset=dataset_encoded["validation"],
    tokenizer=tokenizer
)

trainer.train()

preds_output = trainer.predict(dataset_encoded["validation"])

y_preds = np.argmax(preds_output.predictions, axis=1)
y_valid = np.array(dataset_encoded["validation"]["label"])
labels =['dokujo', 'it', 'sports-','kaden','homme','movie','peachy','smax','topic-news']

def plot_confusion_matrix(y_preds, y_true, labels):
    cm = confusion_matrix(y_true, y_preds, normalize="true")
    fig, ax = plt.subplots(figsize=(6, 6))
    disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=labels)
    disp.plot(cmap="Blues", values_format=".2f", ax=ax, colorbar=False)
    plt.title("Normalized confusion matrix")
    plt.show()

plot_confusion_matrix(y_preds, y_valid,labels)

trainer.save_model('./text/model')