文章分類アノテーションを、ipywidgetsでUIを作成することで効率化する


やったこと

文書分類アノテーションを少しでも楽にすべく、notebook上で動くアノテーション用UIをipywidgetsにて実装しました。
下記は実際の画面となります(livedoorニュースをアノテーションする形を想定)。ボタンの押下により文章の表示、カテゴリラベルの選択、csvへの保存ができるようになっています。

コード

環境はGoogle Colaboratoryです。
text列、category列を持つDataFrameがすでに用意されている前提でのコードとなります。今からアノテーションをするのでcategory列はNoneになっています。

import datetime
import pandas as pd
from IPython.display import HTML, display
from ipywidgets import HBox, VBox, interactive, widgets

# 文章を折り返し表示する設定
# 設定が適用されるのは次回以降に実行されるセルのため、先に実行しておく
def set_css():
  display(HTML('''
  <style>
    pre {
        white-space: pre-wrap;
    }
  </style>
  '''))
get_ipython().events.register('pre_run_cell', set_css)
########################################################################
# ボタン押下時の挙動を設定
########################################################################
def show_text():
    print(df.loc[output_row_box.value, ["text"]].values[0])
    print("-"*50)
    print("現在のカテゴリ", df.loc[output_row_box.value, ["category"]].values[0])
    print("-"*50)

def change_category():
    cat_before = df.loc[output_row_box.value, ["category"]].values[0]
    df.loc[output_row_box.value, "category"] = category_btn.value
    print(f"カテゴリ変更: {cat_before}{category_btn.value}")

def make_csv():
    df.to_csv(output_text_box.value)

    current_japan_time = datetime.datetime.now(
        datetime.timezone(datetime.timedelta(hours=9)))
    save_time = current_japan_time.strftime('%Y/%m/%d %H:%M:%S')
    print(f"保存完了:{save_time}")

########################################################################
# widgets(テキストボックス、各種ボタン)の作成
########################################################################
# 表示する行を入力するボックス
output_row_box = widgets.BoundedIntText(min=0, max=len(df))

# テキスト表示の実行ボタン
show_text_btn = interactive(show_text, {'manual': True, 'manual_name': 'テキストを表示'})
show_text_btn.manual_button.style.button_color = 'lightblue'

# カテゴリの選択肢ボタン
categories = ["トピックニュース", "Sports Watch", "ITライフハック","家電チャンネル", 
    "MOVIE ENTER", "独女通信","エスマックス", "livedoor HOMME", "Peachy"]
category_btn = widgets.ToggleButtons(options=categories, disabled=False)

# カテゴリ変更の実行ボタン
change_category_btn = interactive(change_category, {'manual': True, 'manual_name': 'カテゴリを変更'})
change_category_btn.manual_button.style.button_color = 'lightblue'

# csvのファイル名を設定するテキストボックス
output_text_box = widgets.Text("annotation.csv")

# csv出力の実行ボタン
make_csv_btn = interactive(make_csv, {'manual': True, 'manual_name': 'csvに保存'})
make_csv_btn.manual_button.style.button_color = 'lightblue'

########################################################################
# notebook上にwidgetsを表示
########################################################################
# テキスト選択
print("【テキストを選択】")
display(HBox(children=[output_row_box, show_text_btn]))

# カテゴリ変更
print("\r\n【カテゴリを選択】")
category_btn.layout.flex = '1'
change_category_btn.layout.flex = '1'
display(HBox(children=[category_btn, change_category_btn]))

# csv出力
print("\r\n【csvへ保存】")
display(HBox(children=[output_text_box, make_csv_btn]))

感想

カテゴリ一覧が常に見えているので、Excelなどを使うよりも快適にアノテーションができるようになったかなと思います。
ipywidgetsは直感的に操作できるUIが簡単に作れる点が素晴らしいですね。アイディア次第で色々使い道がありそうだなと感じたため、他での活用方法を考えていきたいと思います。