【TensorFlow】StringLookupの使い方

以下のTweetで紹介したTensorFlow Recommendersのコードを見ていたら、StringLookupクラスという見慣れぬクラスを使っていました。あまり紹介している記事を見たことがないので、Tipsとしてどのようなものか紹介します。

TensorFlowで推薦システムを構築するためのパッケージ「TensorFlow Recommenders」の紹介記事。Kerasでコードをシンプルに書ける。初リリースが一週間前なので、これからに期待。https://t.co/CizEfLHWGK
— Hironsan (@Hironsan13) 2020年9月23日

StringLookupとは？

StringLookupは、文字列を整数値に変換するためのクラスです。一般に、機械学習のモデルはその入力として、整数値や実数値などの数値をとります。たとえば自然言語処理のような分野では、機械学習のモデルに単語を入力する際、単語の文字列ではなく整数値のIDに変換して入力します。StringLookupはそのようなときに使うクラスであり、内部に持っている変換のためのルックアップテーブルを使って実現しています。

あらかじめ用意しておいたボキャブラリをvocab引数に渡してルックアップテーブルを作成後、変換したいデータを渡して呼び出すことで、文字列を数値に変換することができます。

>>> vocab = ["a", "b", "c", "d"]
>>> data = tf.constant([["a", "c", "d"], ["d", "z", "b"]])
>>> layer = StringLookup(vocabulary=vocab)
>>> layer(data)
<tf.Tensor: shape=(2, 3), dtype=int64, numpy=
array([[2, 4, 5],
       [5, 1, 3]])>

上の例では、StringLookupのvocab引数に変換用のボキャブラリを渡していましたが、adaptメソッドを使ってデータから構築することもできます。読み込んだデータからルックアップテーブルを作成する場合には便利な方法でしょう。

>>> data = tf.constant([["a", "c", "d"], ["d", "z", "b"]])
>>> layer = StringLookup()
>>> layer.adapt(data)
>>> layer.get_vocabulary()
['', '[UNK]', 'd', 'z', 'c', 'b', 'a']

ちなみに、以下のように長さが異なる入力をadaptに与えるとValueErrorが発生します。

>>> data = [["a", "c", "d"], ["d", "z"]]
>>> layer = StringLookup()
>>> layer.adapt(data)
ValueError: Can't convert non-rectangular Python sequence to Tensor.

このようなときは、可変長の入力を扱えるRaggedTensorにすると、ValueErrorが発生しなくなります。

>>> data = tf.ragged.constant([["a", "c", "d"], ["d", "z"]])
>>> layer = StringLookup()
>>> layer.adapt(data)
>>> layer.get_vocabulary()
['', '[UNK]', 'd', 'z', 'c', 'a']