如何将TFRECORD文件转换为可读文本数据（知识图谱项目场景）

以下代码完整保留了原项目的词表加载、参数校验、映射逻辑，仅将核心写TFRecord逻辑替换为读取解析逻辑，保证转换结果和原始生成逻辑完全对齐，可直接用于验证TFRecord内的加工数据。

完整修改后代码 #

import numpy as np
import os
import tensorflow as tf
import tqdm
import pdb
import glob
import time
import sys
import re
import argparse
import fastBPE
import platform

use_py3 = platform.python_version()[0] == '3'

parser = argparse.ArgumentParser(description='TensorFlow code for reading TFRecords and output readable text')
parser.add_argument('--tfrecord_file', type=str, required=True,
                                        help='待读取的TFRecords文件路径')
parser.add_argument('--control_code', type=str, required=True,
                                        help='生成TFRecords时使用的控制符，必须存在于词表中')
parser.add_argument('--sequence_len', type=int, required=True,
                                        help='生成TFRecords时使用的序列长度（256或512）')
parser.add_argument('--output_text', type=str, required=True,
                                        help='输出可读文本的保存路径')

args = parser.parse_args()

# 加载词表逻辑和原代码完全一致
vocab = open('../vocab').read().decode(encoding='utf-8').split('\n') if not use_py3 else open('../vocab', encoding='utf-8').read().split('\n')
vocab = list(map(lambda x: x.split(' ')[0], vocab)) + ['<unk>'] + ['\n']
print ('{} 个唯一词'.format(len(vocab)))

if args.control_code not in vocab:
    print('提供的控制符不在词表中')
    print('请更换为词表中存在的控制符')
    sys.exit(1)
    
# 索引和词的映射逻辑和原代码完全一致
word2idx = {u:i for i, u in enumerate(vocab)}
idx2word = np.array(vocab)

# TFRecord解析函数，和原代码写入的feature结构完全对齐
def parse_tfrecord_fn(example):
    feature_description = {
        'input': tf.io.FixedLenFeature([args.sequence_len], tf.int64),
        'output': tf.io.FixedLenFeature([args.sequence_len], tf.int64),
    }
    return tf.io.parse_single_example(example, feature_description)

# 加载并解析TFRecord数据集
raw_dataset = tf.data.TFRecordDataset(args.tfrecord_file)
parsed_dataset = raw_dataset.map(parse_tfrecord_fn)

all_tokens = []
for parsed_record in tqdm.tqdm(parsed_dataset):
    # 读取output序列覆盖所有文本内容
    output_ids = parsed_record['output'].numpy()
    output_tokens = idx2word[output_ids]
    all_tokens.extend(output_tokens.tolist())

# 处理BPE标记生成可读文本
full_text = ' '.join(all_tokens)
# 移除BPE连接符
full_text = full_text.replace('@@ ', '')
# 还原换行符
full_text = full_text.replace(' \n ', '\n').replace(' \n', '\n')

# 写入输出文件
with open(args.output_text, 'w', encoding='utf-8') as f:
    f.write(full_text)

print(f'转换完成，可读文本已保存至：{args.output_text}')

使用说明 #

• 运行命令示例：python tfrecord_to_text.py --tfrecord_file 你的输入文件.tfrecords --control_code 生成时用的控制符 --sequence_len 256 --output_text 输出结果.txt
• 注意事项：
  • 词表../vocab路径需和生成TFRecord时的路径保持一致
  • sequence_len参数必须和生成TFRecord时使用的参数完全一致，否则会解析错误
  • 输出文本自动处理了BPE编码标记，直接得到和原始txt接近的可读格式

内容的提问来源于stack exchange，提问作者Abdul Ahad