NYUAD-ComNets/Racial_Bias_Detection_LLaMa

from transformers import pipeline
from transformers import AutoTokenizer ,AutoModelForCausalLM

text="n Arkansas police officer has been fired after telling a group of African-American men that you don’t belong in my city."
prompt='''Classify the text into 0, 1, and return the answer as the corresponding label.
text: {}
label: '''.format(text)

tokenizer = AutoTokenizer.from_pretrained("NYUAD-ComNets/Racial_Bias_Detection_LLaMa")
tokenizer.pad_token_id = tokenizer.eos_token_id

model = AutoModelForCausalLM.from_pretrained(
    "NYUAD-ComNets/Racial_Bias_Detection_LLaMa",
    device_map="auto",
    torch_dtype="float16",
)

pipe = pipeline(task="text-generation", 
                        model=model, 
                        tokenizer=tokenizer, 
                        max_new_tokens=2, 
                        temperature=0.1)
        
result = pipe(prompt)
answer = result[0]['generated_text'].split("label:")[-1].strip()
print(answer)
if('1' in answer):
    print('This text has racial bias')
else:
    print('no racial bias')