<commit_before># splits data into training and validation sets.

import json
import os
import random
import sys

# percent of images to be in the training set
train_percent = 0.66

base_directory = '/public/jthurst3/MemeCaptcha'
image_directory = base_directory + '/images'
input_directory = base_directory + '/cleaned_data'
output_directory = base_directory + '/tensorflow_input_data'

# gets all files in a directory that end with something
def get_files(direc, end):
	files = os.listdir(direc)
	return [f for f in files if f.endswith(end)]

# splits images up into a set of training ones and a set of validation ones
def split_images(image_dir):
	print("splitting images")
	images = get_files(image_dir, '.jpg')
	images = [i.split('.jpg')[0] for i in images]
	# shuffle the images
	random.shuffle(images)
	cutoff = int(len(images)*train_percent)
	train_images = images[:cutoff]
	val_images = images[cutoff:]
	return set(train_images), set(val_images)

# given a set of images, writes a JSON object to a file with all image-caption pairs that
# contain images in the set, from an original list of JSON files.
def write_JSON(image_set, input_json_files, filename):
	j = {}
	j['memes'] = []
	j['images'] = list(image_set)
	memes = j['memes']
	for jfile in input_json_files:
		with open(os.path.join(input_directory, jfile), 'r') as f:
			# each input JSON file will be in the same format as the
			# output JSON file
			caption_data = json.load(f)
			caption_memes = caption_data['memes']
			for meme in caption_memes:
				if meme['image_id'] in image_set:
					memes.append(meme)
		print("finished reading from", jfile)
	# write it to the output file
	with open(filename, 'w') as f:
		json.dump(j, f)

# main method: partition data set
def partition_dataset():
	# check if either output file exists. Return an error message if so.
	train_file = os.path.join(output_directory, "train.json")
	validation_file = os.path.join(output_directory, "val.json")
	train_images_file = os.path.join(output_directory, "train.images")
	validation_images_file = os.path.join(output_directory, "val.images")
	for f in [train_file, validation_file, train_images_file,
			validation_images_file]:
		if os.path.exists(f):
			print("File ", f, "already exists. Please remove before running.")
			sys.exit(1)
	train_images, val_images = split_images(image_directory)
	input_files = get_files(input_directory, '.json')
	write_JSON(train_images, input_files, train_file)
	write_JSON(val_images, input_files, validation_file)
	# record image partition to files
	with open(train_images_file, 'w') as f:
		for image in train_images:
			f.write(image + '\n')
	with open(validation_images_file, 'w') as f:
		for image in val_images:
			f.write(image + '\n')

if __name__ == '__main__':
	partition_dataset()

<commit_msg>if image partition exists, load it<commit_after># splits data into training and validation sets.

import json
import os
import random
import sys

# percent of images to be in the training set
train_percent = 0.66

base_directory = '/public/jthurst3/MemeCaptcha'
image_directory = base_directory + '/images'
input_directory = base_directory + '/cleaned_data'
output_directory = base_directory + '/tensorflow_input_data'

# gets all files in a directory that end with something
def get_files(direc, end):
	files = os.listdir(direc)
	return [f for f in files if f.endswith(end)]

# splits images up into a set of training ones and a set of validation ones
def split_images(image_dir):
	print("splitting images")
	images = get_files(image_dir, '.jpg')
	images = [i.split('.jpg')[0] for i in images]
	# shuffle the images
	random.shuffle(images)
	cutoff = int(len(images)*train_percent)
	train_images = images[:cutoff]
	val_images = images[cutoff:]
	return set(train_images), set(val_images)

# gets image split from train.images and val.images
def get_image_split(image_dir):
	print("getting existing split from train.images and val.images")
	train_images = set()
	val_images = set()
	with open(os.path.join(image_dir), 'train.images') as f:
		train_images = set([s.split() for s in f])
	with open(os.path.join(image_dir), 'val.images') as f:
		val_images = set([s.split() for s in f])
	return train_images, val_images


# given a set of images, writes a JSON object to a file with all image-caption pairs that
# contain images in the set, from an original list of JSON files.
def write_JSON(image_set, input_json_files, filename):
	j = {}
	j['memes'] = []
	j['images'] = list(image_set)
	memes = j['memes']
	for jfile in input_json_files:
		with open(os.path.join(input_directory, jfile), 'r') as f:
			# each input JSON file will be in the same format as the
			# output JSON file
			caption_data = json.load(f)
			caption_memes = caption_data['memes']
			for meme in caption_memes:
				if meme['image_id'] in image_set:
					memes.append(meme)
		print("finished reading from", jfile)
	# write it to the output file
	with open(filename, 'w') as f:
		json.dump(j, f)

# main method: partition data set
def partition_dataset():
	# check if either output file exists. Return an error message if so.
	train_file = os.path.join(output_directory, "train.json")
	validation_file = os.path.join(output_directory, "val.json")
	train_images_file = os.path.join(output_directory, "train.images")
	validation_images_file = os.path.join(output_directory, "val.images")
	files_exist = False
	for f in [train_file, validation_file, train_images_file,
			validation_images_file]:
		if os.path.exists(f):
			files_exist = True
			# print("File ", f, "already exists. Please remove before running.")
			# sys.exit(1)
	if files_exist:
		train_images, val_images = get_image_split(image_directory)
	else:
		train_images, val_images = split_images(image_directory)
	input_files = get_files(input_directory, '.json')
	write_JSON(train_images, input_files, train_file)
	write_JSON(val_images, input_files, validation_file)
	# record image partition to files
	with open(train_images_file, 'w') as f:
		for image in train_images:
			f.write(image + '\n')
	with open(validation_images_file, 'w') as f:
		for image in val_images:
			f.write(image + '\n')

if __name__ == '__main__':
	partition_dataset()


{'fraction_non_alphanumeric': 0.10700389105058365, 'fraction_numerical': 0.0019455252918287938, 'mean_word_length': 4.9203454894433785, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7057366', 'n_tokens_mistral': 2128, 'n_tokens_neox': 2036, 'n_words': 636}