やみんちゅの備忘録

情報系院生(元工大生)の日々のぼやきや技術的な話など

collections.Counterクラスを使ったスコアの保存

お久しぶりです.やみんちゅです.前回記事の更新を放置して久しいですが,それはそれとして,言語モデルから出力された複数の生成文候補とそれぞれの生成確率をあとでリランキングしやすい形で保存するのにcollections.Counterクラスが便利だと知ったので,備忘録として書いておきたいと思います.

そもそもpythonのcollections.Counterクラスって何だっけ

 詳細な解説はこちら

note.nkmk.me

に譲りますが,例えばリスト内の各要素の個数を調べたいときに

import collections

l = [1, 1, 2, 2, 3, 3, 3]
collections.Counter(l)

このようにリストを渡してcollections.Counterクラスをインスタンスすると,

Counter({1: 2, 2: 2, 3: 3})

辞書形式(ここが大事)で各要素をkeyとしてそれぞれの個数がvalueに格納されます.

Counterを辞書として操作する

 本題ですが,ここで大事なのは

  • Counterの扱い方は辞書とほぼ一緒
  • valueには整数値以外も入れられる

という2点です. そこで試しに下のように空のcollections.Counterクラスを生成し,後からkeyと対応するfloat値の追加を行ってみます.

import collections

c = collections.Counter()
c['a'] = 1.05
print(c)

すると結果は

Counter({'a': 1.05})

となって,しっかり辞書と同じように扱えて,floatも扱えることがわかりましたね. valueがどんな値を受け付けるのか詳細には調べていないのですが,pytorchのtensorvalueに入力することができたので,値を比較できるものなら何でも入れられるかと思います(多分).

Counterクラスを使うと何が嬉しいの?

個人的に嬉しいのは,most_common()関数を使って,簡単にvalueの値順に要素を並べたリストを得られることです. 試しに下のようにコードを書いてみます.

import collections

l = [1, 1, 2, 2, 3, 3, 3]
c = collections.Counter(l)
c.most_common()

すると,下のように(key, value)のタプルがvalueが大きい順に格納されたリストが返ってきます.

# result
[(3, 3), (1, 2), (2, 2)]

またmost_common()関数の引数に整数値nを渡すことで,上位n個までの値を取ってくる,という動作も実現できます.

まとめ

collections.Counterクラスを利用して生成候補と生成確率を保存しておくことと辞書形式と同じ操作ができるので,後から別のスコアを加算してmost_common()で並び替える,という実装が簡単に実現できますよ,というのが結論です.もし必要に迫られたら,皆さんもぜひ使ってみてください.