Update README.md

README.md

@@ -1,3 +1,101 @@
----
-license: mit
----
+---
+license: mit
+base_model:
+- meta-llama/Llama-3.2-1B
+---
+
+## Overview
+A brief description of what this model does and how it’s unique or relevant:
+
+- **Goal**: Classification upon safety of the input text sequences.
+- **Model Description**: DuoGuard-1B-Llama-3.2-transfer is a multilingual, decoder-only LLM-based classifier specifically designed for safety content moderation  across 12 distinct subcategories. Each forward pass produces a 12-dimensional logits vector, where each dimension corresponds to a specific content risk area, such as violent crimes, hate, or sexual content. By applying a sigmoid function to these logits, users obtain a multi-label probability distribution, which allows for fine-grained detection of potentially unsafe or disallowed content.
+For simplified binary moderation tasks, the model can be used to produce a single “safe”/“unsafe” label by taking the maximum of the 12 subcategory probabilities and comparing it to a given threshold (e.g., 0.5). If the maximum probability across all categories is above the threshold, the content is deemed “unsafe.” Otherwise, it is considered “safe.”
+
+DuoGuard-1B-Llama-3.2-transfer is built upon Llama-3.2-1B, a multilingual large language model supporting 8 languages—including English, German, French, Italian, Portuguese, Hindi, Spanish, and Thai. We directly leverage the training data developed fro DuoGuard-0.5B to train Llama-3.2-1B and obtain DuoGuard-1B-Llama-3.2-transfer. Thus, it is specialized (fine-tuned) for safety content moderation primarily in English, French, German, and Spanish, while still retaining the broader language coverage inherited from the Llama-3.2-1B base model. It is provided with open weights.
+## How to Use
+A quick code snippet or set of instructions on how to load and use your model in an application or script:
+```python
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import torch
+
+# 1. Initialize the tokenizer
+tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B")
+tokenizer.pad_token = tokenizer.eos_token
+
+# 2. Load the DuoGuard-0.5B model
+model = AutoModelForSequenceClassification.from_pretrained(
+    "DuoGuard/DuoGuard-1B-Llama-3.2-transfer", 
+    torch_dtype=torch.bfloat16
+).to('cuda:0')
+
+# 3. Define a sample prompt to test
+prompt = "How to kill a python process?"
+
+# 4. Tokenize the prompt
+inputs = tokenizer(
+    prompt,
+    return_tensors="pt", 
+    truncation=True, 
+    max_length=512  # adjust as needed
+).to('cuda:0')
+
+# 5. Run the model (inference)
+with torch.no_grad():
+    outputs = model(**inputs)
+    # DuoGuard outputs a 12-dimensional vector (one probability per subcategory).
+    logits = outputs.logits  # shape: (batch_size, 12)
+    probabilities = torch.sigmoid(logits)  # element-wise sigmoid
+
+# 6. Multi-label predictions (one for each category)
+threshold = 0.5
+category_names = [
+    "Violent crimes",
+    "Non-violent crimes",
+    "Sex-related crimes",
+    "Child sexual exploitation",
+    "Specialized advice",
+    "Privacy",
+    "Intellectual property",
+    "Indiscriminate weapons",
+    "Hate",
+    "Suicide and self-harm",
+    "Sexual content",
+    "Jailbreak prompts",
+]
+
+# Extract probabilities for the single prompt (batch_size = 1)
+prob_vector = probabilities[0].tolist()  # shape: (12,)
+
+predicted_labels = []
+for cat_name, prob in zip(category_names, prob_vector):
+    label = 1 if prob > threshold else 0
+    predicted_labels.append(label)
+
+# 7. Overall binary classification: "safe" vs. "unsafe"
+#    We consider the prompt "unsafe" if ANY category is above the threshold.
+max_prob = max(prob_vector)
+overall_label = 1 if max_prob > threshold else 0  # 1 => unsafe, 0 => safe
+
+# 8. Print results
+print(f"Prompt: {prompt}\n")
+print(f"Multi-label Probabilities (threshold={threshold}):")
+for cat_name, prob, label in zip(category_names, prob_vector, predicted_labels):
+    print(f"  - {cat_name}: {prob:.3f}")
+
+print(f"\nMaximum probability across all categories: {max_prob:.3f}")
+print(f"Overall Prompt Classification => {'UNSAFE' if overall_label == 1 else 'SAFE'}")
+```
+
+### Citation
+
+```plaintext
+@misc{deng2024duoguard,
+      title={DuoGuard: A Two-Player RL-Driven Framework for Multilingual LLM Guardrails}, 
+      author={Yihe Deng and Yu Yang and Junkai Zhang and Wei Wang and Bo Li},
+      year={2024},
+      eprint={2407.},
+      archivePrefix={arXiv},
+      primaryClass={cs.CL},
+      url={https://arxiv.org/abs/2407.}, 
+}
+```