Upload 5 files

LionGuard2.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:20665c1cde68b57c34444accc4f0fca5a3f58b3483d6bad2d6c6911e431afac9
+size 3398496

app.py

@@ -0,0 +1,203 @@
+import gradio as gr
+import openai
+import os
+import sys
+import torch
+
+# # Add the parent directory to the path to import from final_model
+# sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'final_model'))
+
+from safetensors.torch import load_file
+from lionguard2 import LionGuard2
+from utils import get_embeddings
+
+# Set up OpenAI client
+client = openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
+
+# Load LionGuard2 model
+model = LionGuard2()
+model.eval()
+
+# Load model weights
+model_path = os.path.join(os.path.dirname(__file__), '..', 'final_model', 'LionGuard2.safetensors')
+state_dict = load_file(model_path)
+model.load_state_dict(state_dict)
+
+def lionguard_2(message, threshold=0.5):
+    """
+    LionGuard 2 function that uses the actual model to determine if content is unsafe.
+    
+    Args:
+        message: The text message to check
+        threshold: Probability threshold for flagging content as unsafe (default: 0.5)
+    
+    Returns:
+        bool: True if content is flagged as unsafe, False otherwise
+    """
+    try:
+        # Get embeddings for the message
+        embeddings = get_embeddings([message])
+        
+        # Get predictions from the model
+        results = model.predict(embeddings)
+        
+        # Check the binary classification result (overall safety)
+        binary_prob = results['binary'][0]  # First (and only) message's binary probability
+        
+        # Flag as unsafe if probability exceeds threshold
+        return binary_prob > threshold
+        
+    except Exception as e:
+        print(f"Error in LionGuard 2: {e}")
+        # In case of error, default to not flagging to avoid blocking legitimate content
+        return False
+
+def get_openai_response(message, system_prompt="You are a helpful assistant."):
+    """Get response from OpenAI API"""
+    try:
+        response = client.chat.completions.create(
+            model="gpt-4.1-nano",
+            messages=[
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": message}
+            ],
+            max_tokens=500,
+            temperature=0,
+            seed=42,
+        )
+        return response.choices[0].message.content
+    except Exception as e:
+        return f"Error: {str(e)}. Please check your OpenAI API key."
+
+def openai_moderation(message):
+    """
+    OpenAI moderation function that uses OpenAI's built-in moderation API.
+    
+    Args:
+        message: The text message to check
+    
+    Returns:
+        bool: True if content is flagged as unsafe, False otherwise
+    """
+    try:
+        response = client.moderations.create(input=message)
+        return response.results[0].flagged
+    except Exception as e:
+        print(f"Error in OpenAI moderation: {e}")
+        # In case of error, default to not flagging
+        return False
+
+def process_message(message, history_no_mod, history_openai, history_lg):
+    """Process message for all three chatbots"""
+    if not message.strip():
+        return history_no_mod, history_openai, history_lg, ""
+    
+    # Process for gpt-4.1-nano (no moderation)
+    no_mod_response = get_openai_response(message)
+    history_no_mod.append({"role": "user", "content": message})
+    history_no_mod.append({"role": "assistant", "content": no_mod_response})
+    
+    # Process for gpt-4.1-nano with OpenAI moderation
+    openai_flagged = openai_moderation(message)
+    history_openai.append({"role": "user", "content": message})
+    
+    if openai_flagged:
+        openai_response = "🚫 This message has been flagged by OpenAI moderation"
+        history_openai.append({"role": "assistant", "content": openai_response})
+    else:
+        openai_response = get_openai_response(
+            message, 
+        )
+        history_openai.append({"role": "assistant", "content": openai_response})
+    
+    # Process for gpt-4.1-nano with LionGuard 2
+    lg_flagged = lionguard_2(message)
+    history_lg.append({"role": "user", "content": message})
+    
+    if lg_flagged:
+        lg_response = "🚫 This message has been flagged by LionGuard 2"
+        history_lg.append({"role": "assistant", "content": lg_response})
+    else:
+        lg_response = get_openai_response(
+            message, 
+        )
+        history_lg.append({"role": "assistant", "content": lg_response})
+    
+    return history_no_mod, history_openai, history_lg, ""
+
+def clear_all_chats():
+    """Clear all chat histories"""
+    return [], [], []
+
+# Create the Gradio interface
+with gr.Blocks(title="LionGuard 2", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# EMNLP 2025 System Demonstration: LionGuard 2 🦁")
+    gr.Markdown("**LionGuard 2 is a content moderator localised to Singapore - use it to detect unsafe LLM inputs and outputs**")
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            gr.Markdown("## 🔵 No Moderation")
+            chatbot_no_mod = gr.Chatbot(
+                height=800,
+                label="No Moderation",
+                show_label=False,
+                bubble_full_width=False,
+                type='messages'
+            )
+        
+        with gr.Column(scale=1):
+            gr.Markdown("## 🟠 OpenAI Moderation")
+            chatbot_openai = gr.Chatbot(
+                height=800,
+                label="OpenAI Moderation",
+                show_label=False,
+                bubble_full_width=False,
+                type='messages'
+            )
+        
+        with gr.Column(scale=1):
+            gr.Markdown("## 🛡️ LionGuard 2")
+            chatbot_lg = gr.Chatbot(
+                height=800,
+                label="LionGuard 2",
+                show_label=False,
+                bubble_full_width=False,
+                type='messages'
+            )
+    
+    # Single input for all chatbots
+    gr.Markdown("### 💬 Send Message to All Models")
+    with gr.Row():
+        message_input = gr.Textbox(
+            placeholder="Type your message to compare responses...",
+            show_label=False,
+            scale=4
+        )
+        send_btn = gr.Button("Send", variant="primary", scale=1)
+    
+    # Control buttons
+    with gr.Row():
+        clear_btn = gr.Button("Clear All Chats", variant="stop")
+        
+    # Event handlers
+    send_btn.click(
+        process_message,
+        inputs=[message_input, chatbot_no_mod, chatbot_openai, chatbot_lg],
+        outputs=[chatbot_no_mod, chatbot_openai, chatbot_lg, message_input]
+    )
+    
+    message_input.submit(
+        process_message,
+        inputs=[message_input, chatbot_no_mod, chatbot_openai, chatbot_lg],
+        outputs=[chatbot_no_mod, chatbot_openai, chatbot_lg, message_input]
+    )
+    
+    # Clear button
+    clear_btn.click(
+        clear_all_chats,
+        outputs=[chatbot_no_mod, chatbot_openai, chatbot_lg]
+    )
+
+# Launch the app
+if __name__ == "__main__":
+    demo.launch(share=True, debug=True)

lionguard2.py

@@ -0,0 +1,170 @@
+"""
+lionguard2.py
+"""
+
+import torch
+import torch.nn as nn
+
+CATEGORIES = {
+    "binary": ["binary"],
+    "hateful": ["hateful_l1", "hateful_l2"],
+    "insults": ["insults"],
+    "sexual": [
+        "sexual_l1",
+        "sexual_l2",
+    ],
+    "physical_violence": ["physical_violence"],
+    "self_harm": ["self_harm_l1", "self_harm_l2"],
+    "all_other_misconduct": [
+        "all_other_misconduct_l1",
+        "all_other_misconduct_l2",
+    ],
+}
+
+INPUT_DIMENSION = 3072  # length of OpenAI embeddings
+
+
+class LionGuard2(nn.Module):
+    def __init__(
+        self,
+        input_dim=INPUT_DIMENSION,
+        label_names=CATEGORIES.keys(),
+        categories=CATEGORIES,
+    ):
+        """
+        LionGuard2 is a localised content moderation model that flags whether text violates the following categories:
+
+        1. `hateful`: Text that discriminates, criticizes, insults, denounces, or dehumanizes a person or group on the basis of a protected identity.
+
+        There are two sub-categories for the `hateful` category:
+        a. `level_1_discriminatory`: Text that contains derogatory or generalized negative statements targeting a protected group.
+        b. `level_2_hate_speech`: Text that explicitly calls for harm or violence against a protected group; or language praising or justifying violence against them.
+
+        2. `insults`: Text that insults demeans, humiliates, mocks, or belittles a person or group **without** referencing a legally protected trait.
+        For example, this includes personal attacks on attributes such as someone’s appearance, intellect, behavior, or other non-protected characteristics.
+
+        3. `sexual`: Text that depicts or indicates sexual interest, activity, or arousal, using direct or indirect references to body parts, sexual acts, or physical traits.
+        This includes sexual content that may be inappropriate for certain audiences.
+
+        There are two sub-categories for the `sexual` category:
+        a. `level_1_not_appropriate_for_minors`: Text that contains mild-to-moderate sexual content that is generally adult-oriented or potentially unsuitable for those under 16.
+            May include matter-of-fact discussions about sex, sexuality, or sexual preferences.
+        b. `level_2_not_appropriate_for_all_ages`: Text that contains content aimed at adults and considered explicit, graphic, or otherwise inappropriate for a broad audience.
+            May include explicit descriptions of sexual acts, detailed sexual fantasies, or highly sexualized content.
+
+        4. `physical_violence`: Text that includes glorification of violence or threats to inflict physical harm or injury on a person, group, or entity.
+
+        5. `self_harm`: Text that promotes, suggests, or expresses intent to self-harm or commit suicide.
+
+        There are two sub-categories for the `self_harm` category:
+        a. `level_1_self_harm_intent`: Text that expresses suicidal thoughts or self-harm intention; or content encouraging someone to self-harm.
+        b. `level_2_self_harm_action`: Text that describes or indicates ongoing or imminent self-harm behavior.
+
+        6. `all_other_misconduct`: This is a catch-all category for any other unsafe text that does not fit into the other categories.
+        It includes text that seeks or provides information about engaging in misconduct, wrongdoing, or criminal activity, or that threatens to harm,
+        defraud, or exploit others. This includes facilitating illegal acts (under Singapore law) or other forms of socially harmful activity.
+
+        There are two sub-categories for the `all_other_misconduct` category:
+        a. `level_1_not_socially_accepted`: Text that advocates or instructs on unethical/immoral activities that may not necessarily be illegal but are socially condemned.
+        b. `level_2_illegal_activities`: Text that seeks or provides instructions to carry out clearly illegal activities or serious wrongdoing; includes credible threats of severe harm.
+
+        Lastly, there is an additional `binary` category (#7) which flags whether the text is unsafe in general.
+
+        The model takes in as input text, after it has been encoded with OpenAI's `text-embedding-3-small` model.
+
+        The model outputs the probabilities of each category being true.
+
+        ================================
+
+        Args:
+            input_dim: The dimension of the input embeddings. This defaults to 3072, which is the dimension of the embeddings from OpenAI's `text-embedding-3-small` model. This should not be changed.
+            label_names: The names of the labels. This defaults to the keys of the CATEGORIES dictionary. This should not be changed.
+            categories: The categories of the labels. This defaults to the CATEGORIES dictionary. This should not be changed.
+
+        Returns:
+            A LionGuard2 model.
+        """
+        super(LionGuard2, self).__init__()
+        self.label_names = label_names
+        self.n_outputs = len(label_names)
+        self.categories = categories
+
+        # Shared layers
+        self.shared_layers = nn.Sequential(
+            nn.Linear(input_dim, 256),
+            nn.ReLU(),
+            nn.Dropout(0.2),
+            nn.Linear(256, 128),
+            nn.ReLU(),
+            nn.Dropout(0.2),
+        )
+
+        # Output heads for each label
+        self.output_heads = nn.ModuleList(
+            [
+                nn.Sequential(
+                    nn.Linear(128, 32),
+                    nn.ReLU(),
+                    nn.Linear(32, 2),  # 2 thresholds for ordinal classification
+                    nn.Sigmoid(),
+                )
+                for _ in range(self.n_outputs)
+            ]
+        )
+
+    def forward(self, x):
+        # Pass through shared layers
+        h = self.shared_layers(x)
+        # Pass through each output head
+        return [head(h) for head in self.output_heads]
+
+    def predict(self, embeddings):
+        """
+        Predict the probabilities of each label being true.
+
+        Args:
+            embeddings: A numpy array of embeddings (N * INPUT_DIMENSION)
+
+        Returns:
+            A dictionary of probabilities.
+        """
+        # Convert input to PyTorch tensor if not already
+        if not isinstance(embeddings, torch.Tensor):
+            x = torch.tensor(embeddings, dtype=torch.float32)
+        else:
+            x = embeddings
+
+        # Pass through model
+        with torch.no_grad():
+            outputs = self.forward(x)
+
+        # Stack outputs into a single tensor
+        raw_predictions = torch.stack(outputs)  # SIZE:
+
+        # Extract and format probabilities from raw predictions
+        output = {}
+        for i, main_cat in enumerate(self.label_names):
+            sub_categories = self.categories[main_cat]
+            for j, sub_cat in enumerate(sub_categories):
+                # j=0 uses P(y>0)
+                # j=1 uses P(y>1) if L2 category exists
+                output[sub_cat] = raw_predictions[i, :, j]
+
+            # Post processing step:
+            # If L2 category exists, and P(L2) > P(L1),
+            # Set both P(L1) and P(L2) to their average to maintain ordinal consistency
+            if len(sub_categories) > 1:
+                l1 = output[sub_categories[0]]
+                l2 = output[sub_categories[1]]
+
+                # Update probabilities on samples where P(L2) > P(L1)
+                mask = l2 > l1
+                mean_prob = (l1 + l2) / 2
+                l1[mask] = mean_prob[mask]
+                l2[mask] = mean_prob[mask]
+                output[sub_categories[0]] = l1
+                output[sub_categories[1]] = l2
+
+        for key, value in output.items():
+            output[key] = value.numpy().tolist()
+        return output

requirements.txt

@@ -0,0 +1,4 @@
+numpy==2.1.3
+openai==1.83.0
+safetensors==0.5.3
+torch==2.7.0

utils.py

@@ -0,0 +1,44 @@
+"""
+utils.py
+"""
+
+# Standard imports
+import os
+from typing import List
+
+# Third party imports
+import numpy as np
+from openai import OpenAI
+
+client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
+
+# Maximum tokens for text-embedding-3-small
+MAX_TOKENS = 8191  # We don't have access to the tokenizer for text-embedding-3-small, and just assume 1 character = 1 token here
+
+
+def get_embeddings(
+    texts: List[str], model: str = "text-embedding-3-large"
+) -> List[List[float]]:
+    """
+    Generate embeddings for a list of texts using OpenAI API synchronously.
+
+    Args:
+        texts: List of strings to embed.
+        model: OpenAI embedding model to use (default: text-embedding-3-small).
+
+    Returns:
+        A list of embeddings (each embedding is a list of floats).
+
+    Raises:
+        Exception: If the OpenAI API call fails.
+    """
+
+    # Truncate texts to max token limit
+    truncated_texts = [text[:MAX_TOKENS] for text in texts]
+
+    # Make the API call
+    response = client.embeddings.create(input=truncated_texts, model=model)
+
+    # Extract embeddings from response
+    embeddings = np.array([data.embedding for data in response.data])
+    return embeddings