app.py

# This is a project of Chakra Lab LLC. All rights reserved.

import gradio as gr
import os
import time

import torch
import torch.nn.functional as F
from peft import PeftConfig, PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

device = 'cuda' if torch.cuda.is_available() else 'cpu'

#base_model_name = "google/gemma-2b"
#base_model_name = "google/gemma-7b"
base_model_name = "google/gemma-2-9b"
#adapter_model_name = "samidh/cope-g2b-2c-hs.s1.5fpc.9-sx.s1.5.9o-vl.s1.5.9-hr.s5-sh.s5.l5e5-e3-d25-r8"
#adapter_model_name = "samidh/cope-g7bq-2c-hs.s1.5fpc.9-sx.s1.5.9o-VL.s1.5.9-HR.s5-SH.s5-l5e5-e3-d25-r8"

#adapter_model_name = "cope-project/cope-g2.9b-2c-hs.s1.5fpc.9-sx.s1.5.9o-vl.s1.5.9-hr.s5-sh.s5.l5e5-e3-d0-r8"
#adapter_model_name = "samidh/cope-g2.9b-2c-hs.at-sx.at-vl.s5-sh.s5-hr.s5-e2"
#adapter_model_name = "samidh/cope.a.r09.2"
#adapter_model_name = "samidh/cope.a.r09.4.no.vl"
#adapter_model_name = "samidh/cope.a.r09.3"
adapter_model_name = "samidh/cope.a.r09.2"

bnb_config = BitsAndBytesConfig(
    load_in_8bit=True,
)

"""
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True
)
"""

model = AutoModelForCausalLM.from_pretrained(base_model_name,
                                             token=os.environ['HF_TOKEN'],
                                             quantization_config=bnb_config,
                                             device_map="auto")
model = PeftModel.from_pretrained(model, adapter_model_name, token=os.environ['HF_TOKEN'])
model.merge_and_unload()

model = model.to(device)

tokenizer = AutoTokenizer.from_pretrained(base_model_name)

#gr.load("models/samidh/cope-gemma2b-hs-2c-skr-s1.5.9.d25", hf_token=os.environ['HF_TOKEN']).launch()

PROMPT = """
INSTRUCTIONS
============

Examine the given POLICY and determine if the given CONTENT meets the criteria for ANY of the LABELS. Answer "1" if yes, and "0" if no.

POLICY
======

{policy}


CONTENT
=======

{content}


ANSWER
======

"""

DEFAULT_POLICY = """# Policy

## Overview

This policy is designed to determine whether or not content is hate speech.

## Definition of Terms

- Protected Classes: Individuals or groups distinguished by race (e.g. Black, White), ethnicity (e.g. Hispanic), nationality (e.g. Canadian, Argentinian), caste (e.g. Brahmins, Kshatriyas, Vaishyas), immigration status (e.g. Immigrants), religion (e.g. Catholics, Jews), disability (e.g. blind people, deaf people), disease status (e.g. people with AIDS), sexual orientation (e.g. gays, bisexuals), transgender status (e.g. trans people, cis people), gender identity (e.g. women, men), gender expression (e.g. masculine, feminine), or sex (e.g. male, female).
- Derogatory Slur: Any term or phrase used pejoratively or contemptuously towards a member or members of a protected class, including those historically reclaimed but used in a derogatory context.
- Non-Protected Classes: Categories defined by, political party membership (e.g. Democrats, Tories), political beliefs (e.g. fascists, communists), hobbies (e.g. games, knitters), criminal history (e.g. criminals, convicts), or profession (e.g. doctors, farmers).

## Interpretation of Language:

- Terminology or phrasing that can be construed in various ways IS hate speech UNLESS context clearly indicates it unambiguously IS NOT hate speech. So, for example, attacks on a religion as an institution (e.g. "Hinduism is a garbage religion") SHOULD be understood as attacks on the followers of that religion.

## Definition of Labels

### (HS): Hate Speech

#### Includes

- Direct and Indirect Violence Advocacy: Any form of expression that celebrates, justifies, incites, or threatens physical or psychological harm against individuals or groups identified by their membership in a protected class.
- Dehumanization and Vilification: Content suggesting that individuals or groups, explicitly identified by their membership in a specific protected class, are subhuman, inherently lesser beings, or using animalistic descriptions in a way that promotes disdain or hate.
- Derogatory and Dehumanizing Language: Use of slurs, epithets, or any derogatory language aimed at belittling, humiliating, or inciting hatred against individuals or groups explicitly identified by their membership in a specific protected class.
- Explicit and Implicit Discrimination Advocacy: Promoting exclusion, segregation, or denial of rights against  individuals or groups explicitly identified by their membership in a specific protected class.
- Collective Attribution of Negative Actions: Assigning collective blame or advocating collective punishment based on the actions or perceived characteristics of individuals or groups identified by their membership in a protected class.
- Inferiority and Superiority Claims: Statements that categorically assign inferiority or superiority, moral or intellectual, to individuals or groups identified by their membership in a protected class.
- Denial or Distortion of Historical Atrocities: Denying, grossly trivializing, or distorting documented atrocities against groups identified by their membership in a protected class, undermining their significance or the suffering of their members.
- Conspiracy Theories: Propagating unfounded allegations that individuals or groups, identified by their membership in a protected class, are responsible for serious harms or controlling significant institutions to the detriment of society.

#### Excludes

- Attacks on Non-Protected Classes: Content that attacks or criticizes individuals or groups identified by their membership in a Non-Protected Class, EVEN if that attack is violent, threatening, or otherwise hateful (e.g. "Criminals should all be rounded up and shot!").
- Criticism of Beliefs and Institutions: Constructive critique or discussion of political ideologies, religious doctrines, or institutions without resorting to hate speech or targeting individuals or groups identified by their membership in a protected class.
- Neutrally Reporting Historical Events: Neutrally and descriptively reporting or discussion of factual events in the past that could be construed as negative about individuals or groups identified by their membership in a protected class.
- Pushing Back on Hateful Language: Content where the writer pushes back on, condemns, questions, criticizes, or mocks a different person's hateful language or ideas.
- Artistic and Educational Content: Content with legitimate artistic, educational, or documentary value that discusses or portrays issues related to hate speech in a context clearly aimed at enlightening or informing without promoting hate.
"""

DEFAULT_CONTENT = "Put your content sample here."

# Function to make predictions
def predict(content, policy):
    print(f"New request at {time.time()}")

    input_text = PROMPT.format(policy=policy, content=content)
    input_ids = tokenizer.encode(input_text, return_tensors="pt").to(model.device)

    with torch.inference_mode():
        outputs = model(input_ids)

        # Get logits for the last token
        logits = outputs.logits[:, -1, :]  
        
        """
        # Apply softmax to get probabilities
        probabilities = F.softmax(logits, dim=-1)
        
        # Get the predicted token ID
        predicted_token_id = torch.argmax(logits, dim=-1).item()
        
        # Decode the predicted token
        decoded_output = tokenizer.decode([predicted_token_id])
        
        # Get the probability of the predicted token
        predicted_prob = probabilities[0, predicted_token_id].item()
        
        # Function to get probability for a specific token
        def get_token_probability(token):
            token_id = tokenizer.encode(token, add_special_tokens=False)[0]
            return probabilities[0, token_id].item()

        predicted_prob_0 = get_token_probability('0')
        predicted_prob_1 = get_token_probability('1')
        """

        # Get token IDs for "0" and "1"
        token_id_0 = tokenizer.encode("0", add_special_tokens=False)[0]
        token_id_1 = tokenizer.encode("1", add_special_tokens=False)[0]

        # Extract logits for "0" and "1"
        binary_logits = logits[:, [token_id_0, token_id_1]]

        # Apply softmax to get probabilities for these two tokens
        probabilities = F.softmax(binary_logits, dim=-1)

        predicted_prob_0 = probabilities[0,0].item()
        predicted_prob_1 = probabilities[0,1].item()

        decoded_output = '0' if predicted_prob_0 > predicted_prob_1 else '1'
        
    if decoded_output == '1':
        return f'VIOLATING (P: {predicted_prob_1:.2f})'
    else:
        return f'NON-Violating (P: {predicted_prob_0:.2f})'


# Function to make predictions in batches
def predict_batch(contents, policies):
    print(f"New batch request at {time.time()}")

    input_texts = [
        PROMPT.format(policy=policy, content=content)
        for content, policy in zip(contents, policies)
    ]

    input_ids = tokenizer(
        input_texts,
        return_tensors="pt",
        padding=True,
        truncation=True
    ).input_ids.to(model.device)

    with torch.inference_mode():
        outputs = model(input_ids)

    # Get logits for the last tokens
    logits = outputs.logits[:, -1, :]

    # Get token IDs for "0" and "1"
    token_id_0 = tokenizer.encode("0", add_special_tokens=False)[0]
    token_id_1 = tokenizer.encode("1", add_special_tokens=False)[0]

    # Extract logits for "0" and "1"
    binary_logits = logits[:, [token_id_0, token_id_1]]

    # Apply softmax to get probabilities for these two tokens
    probabilities = F.softmax(binary_logits, dim=-1)

    probs_0 = probabilities[:, 0].cpu().numpy()
    probs_1 = probabilities[:, 1].cpu().numpy()

    results = []
    for prob_0, prob_1 in zip(probs_0, probs_1):
        if prob_1 > prob_0:
            output = f'VIOLATING (P: {prob_1:.2f})'
        else:
            output = f'NON-Violating (P: {prob_0:.2f})'
        results.append(output)

    print(results)
    return [results]

# Create Gradio interface
iface = gr.Interface(
    fn=predict_batch,
    inputs=[gr.Textbox(label="Content", lines=2, value=DEFAULT_CONTENT), 
            gr.Textbox(label="Policy", lines=10, value=DEFAULT_POLICY)], 
    outputs=[gr.Textbox(label="Result")],
    batch=True,
    max_batch_size=4,
    title="CoPE Dev (Unstable)",
    description="See if the given content violates your given policy."
)

#iface.queue(
#    default_concurrency_limit=2
#)

# Launch the app
iface.launch()