Upload HarmFormer

README.md

@@ -1,3 +1,94 @@
----
-license: mit
----
+# HarmFormer
+
+HarmFormer is a finetuned `allenai/longformer-base-4096`, which was trained to detect potentially harmful content across 5 different harm categories with three dimensions (Safe, Topical, Toxic) across long text and short text scenarios:
+- H: Hate and Violence
+- IH: Ideological Harm
+- SE: Sexual Harm
+- IL: Illegal Activities
+- SI: Self-Inflicted Harm
+
+We create and define HarmFormer to identify and detect harmful content in text data (especially web pages), which can be used for content moderation, safety checks, and other applications where understanding the nature of text's harmfulness is crucial.
+
+More details about HarmFormer can be found in [our paper - Towards Safer Pretraining: Analyzing and Filtering Harmful Content in Webscale datasets for Responsible LLMs](https://arxiv.org/pdf/2505.02009).
+
+## Model Details
+
+- **Base Model:** allenai/longformer-base-4096
+- **Number of Classes:** 5
+- **Risk Levels per Class:** 3
+- **Max Sequence Length:** 1024
+
+## Usage
+
+```python
+from transformers import AutoTokenizer
+from modeling import HarmFormer
+import torch
+
+# Load the model and tokenizer
+model_path = "themendu/HarmFormer"
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+model = HarmFormer.from_pretrained(model_path)
+
+# Prepare input text
+text = "Your text here"
+inputs = tokenizer(
+    text,
+    add_special_tokens=True,
+    max_length=1024,
+    truncation=True,
+    padding='max_length',
+    return_attention_mask=True,
+    return_tensors='pt',
+)
+
+# Run inference
+with torch.no_grad():
+    outputs = model(**inputs)
+    
+# Process outputs
+logits = torch.stack(outputs, dim=0).permute(1, 0, 2)
+probabilities = torch.softmax(logits, dim=-1)
+predictions = [[[round(prob, 3) for prob in class_probs] for class_probs in sample] for sample in probabilities.cpu().tolist()]
+
+print(predictions)
+```
+
+### Batch Processing
+
+For processing multiple texts at once:
+
+```python
+texts = ["Text 1", "Text 2", "Text 3"]
+inputs = tokenizer(
+    texts,
+    add_special_tokens=True,
+    max_length=1024,
+    truncation=True,
+    padding='max_length',
+    return_attention_mask=True,
+    return_tensors='pt',
+)
+
+with torch.no_grad():
+    outputs = model(**inputs)
+    
+logits = torch.stack(outputs, dim=0).permute(1, 0, 2)
+probabilities = torch.softmax(logits, dim=-1)
+predictions = [[[round(prob, 3) for prob in class_probs] for class_probs in sample] for sample in probabilities.cpu().tolist()]
+```
+
+## Citation
+
+If you use this model in your research, please cite:
+```
+@misc{mendu2025saferpretraininganalyzingfiltering,
+      title={Towards Safer Pretraining: Analyzing and Filtering Harmful Content in Webscale datasets for Responsible LLMs}, 
+      author={Sai Krishna Mendu and Harish Yenala and Aditi Gulati and Shanu Kumar and Parag Agrawal},
+      year={2025},
+      eprint={2505.02009},
+      archivePrefix={arXiv},
+      primaryClass={cs.CL},
+      url={https://arxiv.org/abs/2505.02009}, 
+}
+```

config.json

@@ -0,0 +1,13 @@
+{
+  "model_name": "allenai/longformer-base-4096",
+  "num_classes": 5,
+  "num_risk_levels": 3,
+  "max_length": 1024,
+  "class_map_dict": {
+    "0": "H",
+    "1": "IH",
+    "2": "SE",
+    "3": "IL",
+    "4": "SI"
+  }
+}

modeling.py

@@ -0,0 +1,104 @@
+import torch
+import torch.nn as nn
+import os
+import json
+from transformers import LongformerModel, AutoModel, LongformerTokenizerFast, AutoTokenizer, PreTrainedModel
+
+class HarmFormer(PreTrainedModel):
+    def __init__(self, config):
+        super(HarmFormer, self).__init__(config)
+        self.num_classes = config.num_classes
+        self.num_risk_levels = config.num_risk_levels
+        
+        # Base model
+        self.base_model = AutoModel.from_config(config)
+
+        # Classification heads
+        hidden_size = self.base_model.config.hidden_size
+        
+        self.classifiers = nn.ModuleList([
+            nn.Sequential(
+                nn.Linear(hidden_size, 128),
+                nn.ReLU(),
+                nn.Linear(128, self.num_risk_levels)
+            )
+            for _ in range(self.num_classes)
+        ])
+    
+    def forward(self, input_ids=None, attention_mask=None, **kwargs):
+        outputs = self.base_model(input_ids=input_ids, attention_mask=attention_mask)
+        pooled_output = outputs[1]  # Pooled [CLS] token output
+        
+        # Apply classifiers for each task
+        logits = []
+        for classifier in self.classifiers:
+            logits.append(classifier(pooled_output))
+        
+        return logits
+    
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+        # Load config
+        config_path = os.path.join(pretrained_model_name_or_path, "config.json")
+        if os.path.exists(config_path):
+            with open(config_path, 'r') as f:
+                model_config = json.load(f)
+        else:
+            # Try to load from HF Hub
+            from huggingface_hub import hf_hub_download
+            config_path = hf_hub_download(repo_id=pretrained_model_name_or_path, filename="config.json")
+            with open(config_path, 'r') as f:
+                model_config = json.load(f)
+        
+        # Create base model config
+        from transformers import AutoConfig
+        base_model_name = model_config.get("model_name", "allenai/longformer-base-4096")
+        base_config = AutoConfig.from_pretrained(base_model_name)
+        
+        # Add our custom attributes
+        base_config.num_classes = model_config.get("num_classes", 5)
+        base_config.num_risk_levels = model_config.get("num_risk_levels", 3)
+        base_config.architecture = model_config.get("architecture", "SingleFC")
+        
+        # Create model
+        model = cls(base_config)
+        
+        # Load weights
+        checkpoint_path = os.path.join(pretrained_model_name_or_path, "pytorch_model.bin")
+        if os.path.exists(checkpoint_path):
+            state_dict = torch.load(checkpoint_path, map_location="cpu")
+        else:
+            # Try to load from HF Hub
+            checkpoint_path = hf_hub_download(repo_id=pretrained_model_name_or_path, filename="pytorch_model.bin")
+            state_dict = torch.load(checkpoint_path, map_location="cpu")
+        
+        model.load_state_dict(state_dict)
+        model.eval()
+        
+        return model
+
+def predict_batch(model, tokenizer, texts, batch_size=32):
+    device = next(model.parameters()).device
+    predictions = []
+    
+    # Process in batches to avoid OOM
+    for i in range(0, len(texts), batch_size):
+        batch_texts = texts[i:i+batch_size]
+        inputs = tokenizer(
+            batch_texts,
+            add_special_tokens=True,
+            max_length=1024,
+            truncation=True,
+            padding='max_length',
+            return_attention_mask=True,
+            return_tensors='pt',
+        ).to(device)
+        
+        with torch.no_grad():
+            outputs = model(**inputs)
+            logits = torch.stack(outputs, dim=0).permute(1, 0, 2)  # (batch_size, num_classes, num_risk_levels)
+            probs = torch.softmax(logits, dim=-1)
+            batch_preds = [[[round(prob, 3) for prob in class_probs] for class_probs in sample] for sample in probs.cpu().tolist()]
+            predictions.extend(batch_preds)
+    
+    return predictions

pytorch_model.bin

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

special_tokens_map.json

@@ -0,0 +1,15 @@
+{
+  "bos_token": "<s>",
+  "cls_token": "<s>",
+  "eos_token": "</s>",
+  "mask_token": {
+    "content": "<mask>",
+    "lstrip": true,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": "<pad>",
+  "sep_token": "</s>",
+  "unk_token": "<unk>"
+}

tokenizer_config.json

@@ -0,0 +1,15 @@
+{
+  "add_prefix_space": false,
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": true,
+  "cls_token": "<s>",
+  "eos_token": "</s>",
+  "errors": "replace",
+  "mask_token": "<mask>",
+  "model_max_length": 4096,
+  "pad_token": "<pad>",
+  "sep_token": "</s>",
+  "tokenizer_class": "LongformerTokenizer",
+  "trim_offsets": true,
+  "unk_token": "<unk>"
+}