# 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"

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 guides the classification of content related to expressions that may be sensitive or controversial, aiming to foster a community where open dialogue is encouraged while maintaining respect for all individuals.

## Definition of Terms

- Protected Classes: Individuals or groups of individuals defined by sharing a race, ethnicity, religion, disability, sex, gender identity, or sexual orientation. The scope is narrowed to focus on characteristics commonly recognized in anti-discrimination laws. Religions themselves, as concepts ARE NOT protected - only their members are.
- Non-Protected Classes: Groups defined by other characteristics, such as age, socioeconomic status, and education, which are typically excluded partially or entirely from anti-discrimination laws.

## Interpretation of Language


## Definition of Labels

### (HS): Hate Speech

#### Includes

- Explicit Violence Advocacy: Direct and clear threats of or praise for violence or harm against individuals or groups based on their membership in a particular protected class. Stating an intention to defend against aggression DOES NOT qualify
- Intentional Dehumanization: Statements that clearly depict individuals or groups as inherently ""other"", alien, animalistic, unintelligent, immoral, unclean, or less-than-fully-human based on their membership in a particular protected class in a way that justifies harm or discrimination.
- Targeted Use of Derogatory Slurs: Targeting another person or group of people using a one-word name for a particular protected class that has an inherent negative connotation (e.g. Nigger, Kike, Cunt, Retard). Multi-word terms are never slurs.
- Explicit Discrimination Advocacy: Direct and clear calls for exclusion, segregation, or discrimination against individuals or groups based on their membership in a particular protected class, with a clear intent to promote inequality.
- Direct Hateful Insults: Content that directly addresses another person or group of people the second person (e.g. ""You over there"") and insults them based on their membership in a particular protected class

#### Excludes

- Artistic and Educational Content: Expressions intended for artistic, educational, or documentary purposes that discuss sensitive topics but do not advocate for violence or discrimination against individuals or groups based on their membership in a particular protected class.
- Political and Social Commentary: Commentary on political issues, social issues, and political ideologies that does not directly incite violence or discrimination against individuals or groups based on their membership in a particular protected class.
- Rebutting Hateful Language: Content that rebuts, condemns, questions, criticizes, or mocks a different person's hateful language or ideas OR that insults the person advocating those hateful
- Quoting Hateful Language: Content in which the author quotes someone else's hateful language or ideas while discussing, explaining, or neutrally factually presenting those ideas.
- Describing Sectarian Violence: Content that describes, but does not endorse or praise, violent physical injury against a specifically named race, ethnicity, nationality, sexual orientation, or religious community by another specifically named race, ethnicity, nationality, sexual orientation, or religious community
"""

DEFAULT_CONTENT = "LLMs steal our jobs."

# 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=16,
    title="CoPE Dev (Unstable)",
    description="See if the given content violates your given policy."
)

#iface.queue(
#    default_concurrency_limit=4
#)

# Launch the app
iface.launch()