Added testing FALCON

app.py

@@ -872,7 +872,7 @@ with gr.Blocks(theme=soft, title="Social Bias Testing in Language Models",
                         gen_title = gr.Markdown("### Select Tested Model", visible=True)
 
                         # Tested Model Selection - "openlm-research/open_llama_7b"
-                        tested_model_name = gr.Dropdown( ["bert-base-uncased","bert-large-uncased","gpt2","gpt2-medium","gpt2-large","emilyalsentzer/Bio_ClinicalBERT","microsoft/biogpt","openlm-research/open_llama_3b", "openlm-research/open_llama_7b"], value="bert-base-uncased", 
+                        tested_model_name = gr.Dropdown( ["bert-base-uncased","bert-large-uncased","gpt2","gpt2-medium","gpt2-large","emilyalsentzer/Bio_ClinicalBERT","microsoft/biogpt","openlm-research/open_llama_3b", "openlm-research/open_llama_7b", "tiiuae/falcon-7b"], value="bert-base-uncased", 
                             multiselect=None,
                             interactive=True, 
                             label="Tested Language Model", 

mgr_bias_scoring.py

@@ -20,6 +20,8 @@ from transformers import GPT2LMHeadModel, GPT2Tokenizer
 from transformers import BioGptForCausalLM, BioGptTokenizer
 # LLAMA
 from transformers import LlamaTokenizer, LlamaForCausalLM
+# FALCON
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
 import mgr_sentences as smgr
 import mgr_biases as bmgr
@@ -71,7 +73,18 @@ def _getModel(model_name, device):
                                         offload_folder="offload",
                                         offload_state_dict = True,
                                         device_map='auto')
-    #model.tie_weights()
+  elif "falcon" in model_name.lower():
+    print(f"Getting FALCON model: {model_name}")
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name,
+                                        torch_dtype=torch.bfloat16,
+                                        trust_remote_code=True,
+                                        low_cpu_mem_usage=True, ##
+                                        #use_safetensors=True, ##
+                                        offload_folder="offload",
+                                        offload_state_dict = True,
+                                        device_map='auto')
+  #model.tie_weights()
   if model == None:
     print("Model is empty!!!")
   else:
@@ -634,7 +647,7 @@ def testModelProbability(model_name, model, tokenizer, device):
     print(f"Testing on GPT-2 family model: {model_name}")
     #print(getGPT2Prob(model, tokenizer, "African American man was [MASK]", ["poor","rich"], device, verbose=True))
   elif 'llama' in model_name:
-    print(f"Testing on LLAMA family model: {model_name}")
+    print(f"Testing on LLAMA or FALCON family model: {model_name}")
     #print(getGPT2Prob(model, tokenizer, "African American man was [MASK]", ["poor","rich"], device, verbose=True))
 
 # bias test on one row of a dataframe -> row is one sentence template with target terms
@@ -645,7 +658,8 @@ def checkBias(row, biasProbFunc, model, tokenizer, device, progress, df_len):
   if progress != None:
     progress(row.name/df_len, desc=f"{row['template']}")
 
-  test_res = [1,0] # fail-safe
+  test_res = [0,1]
+  random.shuffle(test_res) # fail-safe
   try:
     test_res, sentences = biasProbFunc(model, tokenizer, row['template'].replace("[T]","[MASK]"), grp_terms, device)
   except ValueError as err:
@@ -740,8 +754,8 @@ def testBiasOnPairs(gen_pairs_df, bias_spec, model_name, model, tokenizer, devic
       gen_pairs_df[['stereotyped','top_term','bottom_term','top_logit','bottom_logit']] = gen_pairs_df.progress_apply(
             biasTestFunc, biasProbFunc=modelGPT2TestFunc, model=model, tokenizer=tokenizer, device=device, progress=progress, df_len=gen_pairs_df.shape[0], axis=1)
 
-    elif 'llama' in model_name.lower():
-      print(f"Testing on LLAMA family model: {model_name}")
+    elif 'llama' in model_name.lower() or 'falcon' in model_name.lower():
+      print(f"Testing on LLAMA or FALCON family model: {model_name}")
       gen_pairs_df[['stereotyped','top_term','bottom_term','top_logit','bottom_logit']] = gen_pairs_df.progress_apply(
             biasTestFunc, biasProbFunc=modelGPT2TestFunc, model=model, tokenizer=tokenizer, device=device, progress=progress, df_len=gen_pairs_df.shape[0], axis=1)