Update app.py

app.py

@@ -1,23 +1,39 @@
 import gradio as gr
-from datasets import load_dataset
+from datasets import load_from_disk
 import numpy as np
 
 
-gender_labels = ['man', 'non-binary', 'woman', 'no_gender_specified', ]
+adjectives = ['ambitious', 'assertive', 'committed', 'compassionate', 'confident', 'considerate', 'decisive', 'determined', 'emotional', 'gentle', 
+              'honest', 'intellectual', 'modest', 'no_adjective', 'outspoken', 'pleasant', 'self-confident', 'sensitive', 'stubborn', 'supportive', 'unreasonable']
 
-ethnicity_labels = ['African-American', 'American_Indian', 'Black', 'Caucasian', 'East_Asian',
-                    'First_Nations', 'Hispanic', 'Indigenous_American', 'Latino', 'Latinx',
-                    'Multiracial', 'Native_American', 'Pacific_Islander', 'South_Asian',
-                    'Southeast_Asian', 'White', 'no_ethnicity_specified']
+professions = ['CEO', 'IT_specialist', 'accountant', 'aerospace_engineer', 'aide', 'air_conditioning_installer', 'architect', 'artist', 'author',
+               'baker', 'bartender', 'bus_driver', 'butcher', 'career_counselor', 'carpenter', 'carpet_installer', 'cashier', 'childcare_worker',
+               'civil_engineer', 'claims_appraiser', 'cleaner', 'clergy', 'clerk', 'coach', 'community_manager', 'compliance_officer', 'computer_programmer',
+               'computer_support_specialist', 'computer_systems_analyst', 'construction_worker', 'cook', 'correctional_officer', 'courier', 'credit_counselor',
+               'customer_service_representative', 'data_entry_keyer', 'dental_assistant', 'dental_hygienist', 'dentist', 'designer', 'detective', 'director',
+               'dishwasher', 'dispatcher', 'doctor', 'drywall_installer', 'electrical_engineer', 'electrician', 'engineer', 'event_planner', 'executive_assistant',
+               'facilities_manager', 'farmer', 'fast_food_worker', 'file_clerk', 'financial_advisor', 'financial_analyst', 'financial_manager', 'firefighter',
+               'fitness_instructor', 'graphic_designer', 'groundskeeper', 'hairdresser', 'head_cook', 'health_technician', 'host', 'hostess', 'industrial_engineer',
+               'insurance_agent', 'interior_designer', 'interviewer', 'inventory_clerk', 'jailer', 'janitor', 'laboratory_technician', 'language_pathologist',
+               'lawyer', 'librarian', 'logistician', 'machinery_mechanic', 'machinist', 'maid', 'manager', 'manicurist', 'market_research_analyst',
+               'marketing_manager', 'massage_therapist', 'mechanic', 'mechanical_engineer', 'medical_records_specialist', 'mental_health_counselor',
+               'metal_worker', 'mover', 'musician', 'network_administrator', 'nurse', 'nursing_assistant', 'nutritionist', 'occupational_therapist',
+               'office_clerk', 'office_worker', 'painter', 'paralegal', 'payroll_clerk', 'pharmacist', 'pharmacy_technician', 'photographer',
+               'physical_therapist', 'pilot', 'plane_mechanic', 'plumber', 'police_officer', 'postal_worker', 'printing_press_operator', 'producer',
+               'psychologist', 'public_relations_specialist', 'purchasing_agent', 'radiologic_technician', 'real_estate_broker', 'receptionist',
+               'repair_worker', 'roofer', 'sales_manager', 'salesperson', 'school_bus_driver', 'scientist', 'security_guard', 'sheet_metal_worker', 'singer',
+               'social_assistant', 'social_worker', 'software_developer', 'stocker', 'supervisor', 'taxi_driver', 'teacher', 'teaching_assistant', 'teller',
+               'therapist', 'tractor_operator', 'truck_driver', 'tutor', 'underwriter', 'veterinarian', 'waiter', 'waitress', 'welder', 'wholesale_buyer', 'writer']
 models = ['DallE', 'SD_14', 'SD_2']
 nos = [1,2,3,4,5,6,7,8,9,10]
-indexes = [768, 1536, 10752]
-ds = load_dataset("tti-bias/identities", split="train")
 
-def get_nearest_64(gender, ethnicity, model, no, index):
+ds = load_from_disk("jobs")
+
+def get_nearest(adjective, profession, model, no):
+    index=768
     df = ds.remove_columns(["image","image_path"]).to_pandas()
     index = np.load(f"indexes/knn_{index}_65.npy")
-    ix = df.loc[(df['ethnicity'] == ethnicity)  & (df['gender'] == gender) & (df['no'] == no) & (df['model'] == model)].index[0]
+    ix = df.loc[(df['adjective'] == adjective)  & (df['profession'] == profession) & (df['no'] == no) & (df['model'] == model)].index[0]
     image = ds.select([index[ix][0]])["image"][0]
     neighbors =  ds.select(index[ix][1:25])
     neighbor_images = neighbors["image"]
@@ -34,14 +50,12 @@ with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
             model = gr.Radio(models, label="Model")
-            index = gr.Radio(indexes, label="Visual vocabulary size")
-            gender = gr.Radio(gender_labels, label="Gender label")
+            adjective = gr.Radio(adjectives, label="Adjective")
         with gr.Column():
-            ethnicity = gr.Radio(ethnicity_labels, label="Ethnicity label")
-            no = gr.Radio(nos, label="Image number")
+            profession = gr.Dropdown(professions, label="Profession")
     button = gr.Button(value="Get nearest neighbors")
     with gr.Row():
         image = gr.Image()
         gallery = gr.Gallery().style(grid=4)
-    button.click(get_nearest_64, inputs=[gender, ethnicity, model, no, index], outputs=[image, gallery])
+    button.click(get_nearest, inputs=[adjective, profession, model, no], outputs=[image, gallery])
 demo.launch()