<jupyter_start><jupyter_text>Create Training Data (TFRecords)In order to train efficiently, the data created in Spark needs to be transformed into TFRecords. This code is run after the "training_data_create_SPARK.ipynb" notebook is run.<jupyter_code># !pip install tensorflow==2.4.1
# !pip install transformers
# !pip install pyarrow
import tensorflow as tf
import pandas as pd
import numpy as np
import os
from math import ceil
def check_targets(targs):
    if targs[0] == -1:
        return 1
    else:
        return 0
def create_tfrecords_dataset(data, iter_num, dataset_type='train'):
    data['no_target'] = data['target_tok'].apply(check_targets)
    data = data[data['no_target']==0].copy()
    
    paper_title = tf.ragged.constant(data['paper_title_tok'].to_list())
    
    targets = tf.keras.preprocessing.sequence.pad_sequences(data['target_tok'].to_list(), maxlen=20, 
                                                            dtype='int64', padding='post', 
                                                            truncating='post', value=0)

    ds = tf.data.Dataset.zip((tf.data.Dataset.from_tensor_slices(paper_title),
                              tf.data.Dataset.from_tensor_slices(data['journal_tok'].to_list()),
                              tf.data.Dataset.from_tensor_slices(data['doc_type_tok'].to_list()),
                              tf.data.Dataset.from_tensor_slices(targets)))
    
    serialized_features_dataset = ds.map(tf_serialize_example)
    
    filename = f"./iteration_final/basic_word_tokenized/tfrecords/{dataset_type}/{str(iter_num).zfill(4)}.tfrecord"
    writer = tf.data.experimental.TFRecordWriter(filename)
    writer.write(serialized_features_dataset)
def tf_serialize_example(f0, f1, f2, f3):
    tf_string = tf.py_function(serialize_example, (f0, f1, f2, f3), tf.string)
    return tf.reshape(tf_string, ())
def serialize_example(paper_title, journal, doc_type, targets):
    paper_title_list = tf.train.Int64List(value=paper_title.numpy().tolist())
    journal_list = tf.train.Int64List(value=journal.numpy().tolist())
    doc_type_list = tf.train.Int64List(value=doc_type.numpy().tolist())
    targets_list = tf.train.Int64List(value=targets.numpy().tolist())
    
    paper_title_feature = tf.train.Feature(int64_list = paper_title_list)
    journal_feature = tf.train.Feature(int64_list = journal_list)
    doc_type_feature = tf.train.Feature(int64_list = doc_type_list)
    targets_feature = tf.train.Feature(int64_list = targets_list)
    
    features_for_example = {
        'paper_title': paper_title_feature,
        'journal': journal_feature,
        'doc_type': doc_type_feature,
        'targets': targets_feature
    }
    
    example_proto = tf.train.Example(features=tf.train.Features(feature=features_for_example))
    
    return example_proto.SerializeToString()
def turn_part_file_into_tfrecord(base_path, dataset_type='train'):
    file_list = [x for x in os.listdir(f"{base_path}{dataset_type}") if x.endswith('parquet')]
    file_list.sort()
    print(f"There are {len(file_list)} files for {dataset_type}")
    for i, file_name in enumerate(file_list):
        data = pd.read_parquet(f"{base_path}{dataset_type}/{file_name}")
        print(f"_____File number: {i} ({data.shape[0]} samples)")
        create_tfrecords_dataset(data, i, dataset_type)
base_file_path = f"./iteration_final/tokenized_data/"<jupyter_output><empty_output><jupyter_text>Without padding<jupyter_code>%%time
turn_part_file_into_tfrecord(base_file_path, 'train')
%%time
turn_part_file_into_tfrecord(base_file_path, 'val')
%%time
turn_part_file_into_tfrecord(base_file_path, 'test')<jupyter_output>There are 5 files for test
_____File number: 0 (20395 samples)
_____File number: 1 (24338 samples)
_____File number: 2 (24756 samples)
_____File number: 3 (17072 samples)
_____File number: 4 (10677 samples)
CPU times: user 32.5 s, sys: 2.1 s, total: 34.6 s
Wall time: 28.7 s
{'fraction_non_alphanumeric': 0.11560988031576266, 'fraction_numerical': 0.020626432391138275, 'mean_word_length': 3.9038701622971286, 'pattern_counts': {'":': 0, '<': 8, '<?xml version=': 0, '>': 8, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '172226', 'license': 'MIT', 'n_tokens_mistral': 1334, 'n_tokens_neox': 1286, 'n_words': 236, 'path': 'V1/002_Model/002_data_create_to_tfrecords.ipynb', 'repo_name': 'ourresearch/openalex-concept-tagging', 'size': '3927'}