如何使用Dataset API编写字数统计？

import tensorflow as tf

# 假设有一个包含多个文本文件的文件夹，每个文件包含一段文本
file_pattern = "/path/to/text/files/*.txt"
dataset = tf.data.TextLineDataset(file_pattern)

def count_words(text):
    words = tf.strings.split(text)
    return tf.size(words)

# 对每个文本应用数据处理函数
dataset = dataset.map(count_words)

# 计算总字数
total_words = dataset.reduce(tf.constant(0, dtype=tf.int64), lambda x, y: x + y)

# 打印结果
with tf.Session() as sess:
    print(sess.run(total_words))