DistilRoBERTa-nsfw-prompt-stable-diffusion
This model utilizes the Distilroberta base architecture, which has been fine-tuned for a classification task on AdamCodd/Civitai-2m-prompts dataset, on the positive prompts. It achieves the following results on the evaluation set:
- Loss: 0.3103
- Accuracy: 0.8642
- F1: 0.8612
- AUC: 0.9408
- Precision: 0.8805
- Recall: 0.8427
Model description
This model is designed to identify NSFW prompts in Stable-diffusion, trained on a dataset comprising of ~2 million prompts, evenly split between SFW and NSFW categories (1,043,475 samples of each, ensuring a balanced dataset). Single-word prompts have been excluded to enhance the accuracy and relevance of the predictions.
Additionally, it is important to note that the model assesses the likelihood of a prompt being NSFW based on statistical occurrences, rather than evaluating the specific words. This approach allows for the identification of NSFW content in prompts that may appear SFW. The accuracy of the model tends to increase with the length of the prompt. Therefore, prompts that are extremely brief, such as those comprising only two or three words, might be subject to less accurate evaluations.
Although this model demonstrates satisfactory accuracy, it is recommended to use with this image NSFW detector to improve overall detection capabilities and minimize the occurrence of false positives.
Usage
from transformers import pipeline
prompt_detector = pipeline("text-classification", model="AdamCodd/distilroberta-nsfw-prompt-stable-diffusion")
predicted_class = prompt_detector("masterpiece, 1girl, looking at viewer, sitting, tea, table, garden")
print(predicted_class)
#[{'label': 'SFW', 'score': 0.868}]
Training and evaluation data
More information needed
Training procedure
Training hyperparameters
The following hyperparameters were used during training:
- learning_rate: 3e-05
- train_batch_size: 32
- eval_batch_size: 32
- seed: 42
- optimizer: AdamW with betas=(0.9,0.999) and epsilon=1e-08
- lr_scheduler_type: linear
- lr_scheduler_warmup_steps: 150
- num_epochs: 1
- weight_decay: 0.01
Training results
Metrics: Accuracy, F1, Precision, Recall, AUC
'eval_loss': 0.3103,
'eval_accuracy': 0.8642,
'eval_f1': 0.8612,
'eval_precision': 0.8805,
'eval_recall': 0.8427,
'eval_roc_auc': 0.9408,
[[184931 23859]
[32820 175780]]
Framework versions
- Transformers 4.36.2
- Datasets 2.16.1
- Tokenizers 0.15.0
- Evaluate 0.4.1
If you want to support me, you can here.
- Downloads last month
- 2,404
Finetuned from
Dataset used to train AdamCodd/distilroberta-nsfw-prompt-stable-diffusion
Collection including AdamCodd/distilroberta-nsfw-prompt-stable-diffusion
Evaluation results
- lossself-reported0.310
- Accuracyself-reported0.864
- F1self-reported0.861
- Precisionself-reported0.880
- Recallself-reported0.843
- AUCself-reported0.941