Casting a string tensor to a list of string

Viewed 1468

How could someone cast the following string tensorflow tensor:

<tf.Tensor: shape=(64,), dtype=string, numpy=
array([b'example string 1',
       b'example string 2',
       b'example string 3',

       ...

       b'example string 63',
       b'example string 64'],
      dtype=object)>

to a list of string:

[
    'example string 1',
    'example string 2',
    'example string 3',

       ...

    'example string 63',
    'example string 64'
]

applying a vectorized approach?

Thank you in advance.

3 Answers

Just use list(tensor.numpy()). Example:

import tensorflow as tf

n_strings = 8
t = tf.convert_to_tensor(['example string ' + str(i) for i in range(n_strings)])
t
# <tf.Tensor: shape=(8,), dtype=string, numpy=
# array([b'example string 0', b'example string 1', b'example string 2',
#        b'example string 3', b'example string 4', b'example string 5',
#        b'example string 6', b'example string 7'], dtype=object)>
list(t.numpy().astype('str'))
# ['example string 0',
#  'example string 1',
#  'example string 2',
#  'example string 3',
#  'example string 4',
#  'example string 5',
#  'example string 6',
#  'example string 7']

As an alternative, list(t.numpy().decode('utf-8')) worked for me.

Below you can find a robust conversion function together with unit tests. The function can work with any rank of input string tensor (scalar, 1D, 2D, etc.) and can be parametrized with desired encoding:

import operator
from typing import Any

import numpy as np
import pytest
import tensorflow as tf
from numpy import typing as npt


def parse_string_tensor(tensor: tf.Tensor, encoding: str = "utf-8") -> Any:
    """
    Parse any rank string tensor into a pure-python structure of strings in given encoding.
    """
    # validate that the tensor is a string tensor
    assert tensor.dtype == tf.string
    # convert the tensor to a numpy array of bytes objects
    array: npt.NDArray[np.object_] = tensor.numpy()
    # define the decoding operation
    decoder = np.vectorize(
        operator.methodcaller("decode", encoding), otypes=[np.object_]
    )
    # decode the bytes objects into strings
    decoded = decoder(array)
    # convert to pure-python structure
    return decoded.tolist()


def test_parse_string_tensor_1d() -> None:
    tensor = tf.constant(["", "‍", ""], dtype=tf.string)
    result = parse_string_tensor(tensor, encoding="utf-8")
    assert result == ["", "‍", ""]


def test_parse_string_tensor_wrong_type() -> None:
    tensor = tf.constant([1.7, -0.5, -3.3], dtype=tf.float32)
    with pytest.raises(AssertionError):
        parse_string_tensor(tensor)


def test_parse_string_tensor_scalar() -> None:
    tensor = tf.constant("", dtype=tf.string)
    result = parse_string_tensor(tensor, encoding="utf-8")
    assert result == ""


def test_parse_string_tensor_2d() -> None:
    tensor = tf.constant(
        [["", "‍"], ["", ""]],
        dtype=tf.string,
    )
    result = parse_string_tensor(tensor, encoding="utf-8")
    assert result == [["", "‍"], ["", ""]]


def test_parse_string_tensor_different_encoding() -> None:
    tensor = tf.constant("naïve".encode("latin1"), dtype=tf.string)
    result = parse_string_tensor(tensor, encoding="latin1")
    assert result == "naïve"

def test_parse_string_empty_tensor() -> None:
    tensor = tf.constant([], dtype=tf.string)
    result = parse_string_tensor(tensor, encoding="utf-8")
    assert result == []
Related