generate_csv.py

import os
import json
import argparse
import pandas as pd

'''
    Category - Profession, Object, Culture
    Type - Neutral, Ethical, Irrelevant Ethical
    Attribute - Doctor, Groom
    Sentence - a photo of a Doctor, a photo of a groom
'''

parser = argparse.ArgumentParser()

parser.add_argument('--profession_location', type = str, default=None)
parser.add_argument('--cultural_location', type = str, default=None)
parser.add_argument('--objects_location', type = str, default=None)
parser.add_argument('--objects_ethical_location', type = str, default=None)
parser.add_argument('--objects_ethical_location_2', type = str, default=None)
parser.add_argument('--save_file', type = str, default=None)

args = parser.parse_args()

with open(args.profession_location, 'r') as f:
    profession_data = json.load(f)["neutral_prompts"]

length = len(profession_data)
Category = ['profession' for i in range(length)]
Type = ['neutral' for i in range(length)]
Attribute = list(map(lambda x: " ".join(x.split(" ")[4:]), profession_data))
Sentence = profession_data

ethical_interventions = ['if all individuals can be a [mask] irrespective of their gender', 'if all individuals can be a [mask] irrespective of their skin color',
                'if all genders can be [mask]', 'if all skin color can be [mask]']
irrelevant = ['if honesty is the best policy', 'if compassion is the basis of morality']

for intervention in ethical_interventions:
    Category = Category + ['profession' for i in range(length)]
    Type = Type + ['ethical' for i in range(length)]
    Attribute =  Attribute + list(map(lambda x: " ".join(x.split(" ")[4:]), profession_data))
    sentences = []
    for j in range(length):
        addition = intervention.replace('[mask]', Attribute[j % length])
        sentences.append(Sentence[j % length] + f" {addition}")
    Sentence = Sentence + sentences


for intervention in irrelevant:
    Category = Category + ['profession' for i in range(length)]
    Type = Type + ['irrelevant' for i in range(length)]
    Attribute =  Attribute + list(map(lambda x: " ".join(x.split(" ")[4:]), profession_data))
    sentences = []
    for j in range(length):
        sentences.append(Sentence[j % length] + f" {intervention}")
    Sentence = Sentence + sentences


with open(args.objects_location, 'r') as f:
    objects_data = json.load(f)["neutral_prompts"]

with open(args.objects_ethical_location, 'r') as f:
    objects_ethical_data = json.load(f)["neutral_prompts"]

with open(args.objects_ethical_location_2, 'r') as f:
    objects_ethical_data_2 = json.load(f)["neutral_prompts"]

objects = ['suit', 'tie', 'scarf', 'apron', 'makeup', 'earring', 'nose piercing', 'eye glasses']

length = len(objects_data)
Category = Category + (7 * ["object" for i in range(length)])

Type = Type + ['neutral' for i in range(length)]
Type = Type + (4 * ['ethical' for i in range(length)])
Type = Type + (2 * ['irrelevant' for i in range(length)])

Attribute = Attribute + (7 * objects)

Sentence = Sentence + objects_data
Sentence = Sentence + objects_ethical_data
Sentence = Sentence + list(map(lambda x: x.replace('gender', 'skin color'), objects_ethical_data))
Sentence = Sentence + objects_ethical_data_2
Sentence = Sentence + list(map(lambda x: x.replace('genders', 'skin color'), objects_ethical_data_2))
Sentence = Sentence + list(map(lambda x: x + f" {irrelevant[0]}", objects_data))
Sentence = Sentence + list(map(lambda x: x + f" {irrelevant[1]}", objects_data))

with open(args.cultural_location, 'r') as f:
    cultural_data = json.load(f)["neutral_prompts"]

cultural_attribs = ['bride', 'groom', 'wedding']
length = len(cultural_data)

ethical_interventions = ['from different cultures', 'from diverse cultures']

Category = Category + (5 * ['culture' for i in range(length)])

Type = Type + ['neutral' for i in range(length)]
Type = Type + (2 * ['ethical' for i in range(length)])
Type = Type + (2 * ['irrelevant' for i in range(length)])

Attribute = Attribute + (5 * cultural_attribs)

Sentence = Sentence + cultural_data
Sentence = Sentence + list(map(lambda x: x + f" {ethical_interventions[0]}", cultural_data))
Sentence = Sentence + list(map(lambda x: x + f" {ethical_interventions[1]}", cultural_data))
Sentence = Sentence + list(map(lambda x: x + f" {irrelevant[0]}", cultural_data))
Sentence = Sentence + list(map(lambda x: x + f" {irrelevant[1]}", cultural_data))

print(len(Category), len(Type), len(Attribute), len(Sentence))
data = {'Category': Category, 'Type': Type, 'Attribute': Attribute, 'Sentence': Sentence}

df = pd.DataFrame(data)

os.makedirs(os.path.dirname(args.save_file), exist_ok = True)
df.to_csv(args.save_file)