def make_tfrecord_files()

in recommended-item-search/data_preparation.py [0:0]

• 11 lines of code
• 3 McCabe index (conditional complexity)

def make_tfrecord_files(dataframe, file_type, num_files):
  """Writes training and test data in TFRecord format."""
  user_id = dataframe['user_id'].values
  movie_ids = dataframe['movie_id2'].values
  
  N = dataframe.shape[0]
  for file_id, head_tail in enumerate(split_range(N, num_files)):
    export_file = "{}-{:05d}.tfrecord".format(file_type, file_id)
    export_path = os.path.join(FLAGS.export_dir, export_file)
    with tf.python_io.TFRecordWriter(export_path) as record_writer:
      for i in range(head_tail[0], head_tail[1]):
        sequence_example = make_sequence_example(user_id[i], movie_ids[i])
        record_writer.write(sequence_example.SerializeToString())