python-3.x - tf.py_function からテンソルの辞書を返す方法は?

Question

通常、Transformers トークナイザーは入力を辞書としてエンコードします。

{"input_ids": tf.int32, "attention_mask": tf.int32, "token_type_ids": tf.int32}

また、大規模なデータセットでより優れたパフォーマンス処理をアーカイブするにDataset.mapは、入力データセットの各要素にトークナイザー関数を適用するための使用を含むパイプラインを実装することをお勧めします。Tensorflow チュートリアルで行ったのとまったく同じ: Load text。

ただし、tf.py_function(map python 関数をラップするために使用される) は、上記のようにテンソルの辞書を返すことをサポートしていません。

たとえば、Load テキストのトークナイザー (エンコーダー)が次の辞書を返す場合:

{
    "input_ids": [ 101, 13366,  2131,  1035,  6819,  2094,  1035,  102 ],
    "attention_mask": [ 1, 1, 1, 1, 1, 1, 1, 1 ]
}

テンソルの目的の辞書を取得するためにTout、誰かがのパラメータをどのように設定できますか:tf.py_function

{
    'input_ids': <tf.Tensor: shape=(16,), dtype=int32, numpy = array(
    [ 101, 13366,  2131,  1035,  6819,  2094,  1035,  102 ], dtype=int32)>

    'attention_mask': <tf.Tensor: shape=(16,), dtype=int32, numpy=array(
     [ 1, 1, 1, 1, 1, 1, 1, 1 ], dtype=int32)>
}

?