Added benchmarks

README.md

@@ -2,7 +2,6 @@
 library_name: transformers
 language:
 - en
-pipeline_tag: zero-shot-classification
 ---
 
 # Tiny-Toxic-Detector
@@ -12,14 +11,93 @@ A tiny comment toxicity classifier model at only 2M parameters. With only ~10MB
 A paper on this model is being released soon.
 
 
+### Benchmarks
+
+The Tiny-Toxic-Detector achieves an impressive 90.26% on the Toxigen benchmark and 87.34% on the Jigsaw-Toxic-Comment-Classification-Challenge. Here we compare our results against other toxic classification models:
+
+
+| Model                             | Size (parameters) | Toxigen (%) | Jigsaw (%) |
+| --------------------------------- | ----------------- | ----------- | ---------- |
+| lmsys/toxicchat-t5-large-v1.0     | 738M              | 72.67       | 88.82      |
+| s-nlp/roberta toxicity classifier | 124M              | 88.41       | 94.92      |
+| unitary/toxic-bert                | 109M              | 49.50       | 89.70      |
+| mohsenfayyaz/toxicity-classifier  | 109M              | 81.50       | 83.31      |
+| martin-ha/toxic-comment-model     | 67M               | 68.02       | 91.56      |
+| Tiny-toxic-detector               | 2M                | 90.26       | 87.34      |
+
+
 ### Usage
-This model uses custom architecture and requires some extra custom code to work. Below you can find a script that is fully-usable.
+This model uses custom architecture and requires some extra custom code to work. Below you can find the architecture and a fully-usable example.
+
+<details>
+  <summary>
+    Architecture  
+  </summary>
+
+```python
+import torch
+import torch.nn as nn
+from transformers import PreTrainedModel, PretrainedConfig, AutoTokenizer
+
+# Define TinyTransformer model
+class TinyTransformer(nn.Module):
+    def __init__(self, vocab_size, embed_dim, num_heads, ff_dim, num_layers):
+        super().__init__()
+        self.embedding = nn.Embedding(vocab_size, embed_dim)
+        self.pos_encoding = nn.Parameter(torch.zeros(1, 512, embed_dim))
+        encoder_layer = nn.TransformerEncoderLayer(d_model=embed_dim, nhead=num_heads, dim_feedforward=ff_dim, batch_first=True)
+        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
+        self.fc = nn.Linear(embed_dim, 1)
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x):
+        x = self.embedding(x) + self.pos_encoding[:, :x.size(1), :]
+        x = self.transformer(x)
+        x = x.mean(dim=1)  # Global average pooling
+        x = self.fc(x)
+        return self.sigmoid(x)
+
+class TinyTransformerConfig(PretrainedConfig):
+    model_type = "tiny_transformer"
+
+    def __init__(self, vocab_size=30522, embed_dim=64, num_heads=2, ff_dim=128, num_layers=4, max_position_embeddings=512, **kwargs):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.ff_dim = ff_dim
+        self.num_layers = num_layers
+        self.max_position_embeddings = max_position_embeddings
+
+class TinyTransformerForSequenceClassification(PreTrainedModel):
+    config_class = TinyTransformerConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = 1
+        self.transformer = TinyTransformer(
+            config.vocab_size,
+            config.embed_dim,
+            config.num_heads,
+            config.ff_dim,
+            config.num_layers
+        )
+
+    def forward(self, input_ids, attention_mask=None):
+        outputs = self.transformer(input_ids)
+        return {"logits": outputs}
+```
+</details>
+
+<details>
+  <summary>
+    Full example  
+  </summary>
+
 ```python
 import torch
 import torch.nn as nn
 from transformers import PreTrainedModel, PretrainedConfig, AutoTokenizer
-from huggingface_hub import login
-import os
 
 # Define TinyTransformer model
 class TinyTransformer(nn.Module):
@@ -112,6 +190,16 @@ def main():
 if __name__ == "__main__":
     main()
 ```
+</details>
 
-
-
+Please note that to predict toxicity you can use the following example:
+```python
+inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=128, padding="max_length").to(device)
+if "token_type_ids" in inputs:
+  del inputs["token_type_ids"]
+
+with torch.no_grad():
+  outputs = model(**inputs)
+  logits = outputs["logits"].squeeze()
+  prediction = "Toxic" if logits > 0.5 else "Not Toxic"
+```