simocorbo/toxicthesis-gemini-3.5-flash-roberta-classification-4

RoBERTa - GEMINI-3.5-FLASH - Classification (4 classes)

Toxicity prediction model trained on the GEMINI-3.5-FLASH dataset.

Property	Value
Model	RoBERTa
Task	Classification (4 classes)
Dataset	gemini-3.5-flash
Framework	PyTorch / PyTorch Lightning

Class: RobertaModel

from src.models.roberta import RobertaModel

model = RobertaModel(
    model_name: str = 'roberta-base',  # HuggingFace model name
    num_classes: int = 2,               # 1=regression, 2=binary, 3+=multi-class
    loss_fn: str = 'auto',
    freeze_layers: int = 0,             # Number of layers to freeze
    dropout: float = 0.1,
    lr: float = 2e-5,
    gradient_checkpointing: bool = True
)

Methods

Method	Description
forward(inputs)	Input: dict with 'input_ids', 'attention_mask'. Returns raw logits.
predict_text(text)	Easiest method: Predict on raw text string. Returns dict with score/probabilities.
compute_loss(logits, targets)	Computes loss based on num_classes.
from_config(config)	Create model from config dict.

Usage with ToxicThesis (Recommended)

# 1. Clone ToxicThesis repository
# git clone https://github.com/simo-corbo/ToxicThesis
# cd ToxicThesis && pip install -r requirements.txt

from huggingface_hub import hf_hub_download
import torch

# 2. Download checkpoint
checkpoint_path = hf_hub_download(
    repo_id="simocorbo/toxicthesis-gemini-3.5-flash-roberta-classification-4",
    filename="checkpoints/best.pt"
)

# 3. Import and load model from ToxicThesis
from src.models.roberta import RobertaModel

checkpoint = torch.load(checkpoint_path, map_location='cpu', weights_only=False)
hparams = checkpoint.get('hyper_parameters', {})

model = RobertaModel(
    model_name=hparams.get('model_name', 'roberta-base'),
    num_classes=hparams.get('num_classes', 4),
    dropout=hparams.get('dropout', 0.1)
)
model.load_state_dict(checkpoint.get('state_dict', checkpoint), strict=False)
model.eval()

# 4. Use the built-in predict_text method (easiest!)
result = model.predict_text("Your text here")
print(result)
# Output: {'score': 0.7234, 'prediction': 'Toxicity score: 0.7234', ...}

# 5. Multiple predictions
texts = ["Hello friend", "You are terrible", "Have a nice day"]
for text in texts:
    result = model.predict_text(text)
    print(f"{text}: {result['score']:.4f}")

Standalone Usage (without ToxicThesis)

from huggingface_hub import hf_hub_download
import torch
import torch.nn as nn
from transformers import AutoModel, AutoTokenizer

# 1. Download checkpoint
checkpoint_path = hf_hub_download(
    repo_id="simocorbo/toxicthesis-gemini-3.5-flash-roberta-classification-4",
    filename="checkpoints/best.pt"
)

# 2. Load checkpoint
checkpoint = torch.load(checkpoint_path, map_location='cpu', weights_only=False)
hparams = checkpoint.get('hyper_parameters', {})
model_name = hparams.get('model_name', 'roberta-base')
num_classes = hparams.get('num_classes', 4)

# 3. Define model class
class RobertaClassifier(nn.Module):
    def __init__(self, model_name, num_classes, dropout=0.1):
        super().__init__()
        self.num_classes = num_classes
        self.roberta = AutoModel.from_pretrained(model_name, add_pooling_layer=False)
        hidden_size = self.roberta.config.hidden_size
        self.dropout = nn.Dropout(dropout)
        self.classifier = nn.Linear(hidden_size, 1 if num_classes <= 2 else num_classes)
    
    def forward(self, input_ids, attention_mask=None):
        outputs = self.roberta(input_ids=input_ids, attention_mask=attention_mask)
        pooled = self.dropout(outputs.last_hidden_state[:, 0, :])
        return self.classifier(pooled)

# 4. Load model
model = RobertaClassifier(model_name, num_classes)
state_dict = checkpoint.get('state_dict', checkpoint)
model.load_state_dict(state_dict, strict=False)
model.eval()

# 5. Inference
tokenizer = AutoTokenizer.from_pretrained(model_name)

def predict(text: str) -> dict:
    inputs = tokenizer(text, return_tensors='pt', truncation=True, max_length=512)
    with torch.no_grad():
        logits = model(inputs['input_ids'], inputs.get('attention_mask'))
        if num_classes == 1:
            return {'score': torch.sigmoid(logits).item()}
        elif num_classes == 2:
            p = torch.sigmoid(logits).item()
            return {'probability': p, 'class': int(p >= 0.5)}
        else:
            probs = torch.softmax(logits, dim=-1).squeeze().tolist()
            return {'probabilities': probs, 'class': int(torch.argmax(torch.tensor(probs)))}

result = predict("Your text here")
print(result)

Score Interpretation

Output	Range	Meaning
probabilities	List[float]	Probability distribution over 4 classes.
class	0 to 3	Predicted class (argmax of probabilities).

Classes: 4 toxicity levels, where higher class index = more toxic.

Files

File	Description
checkpoints/best.pt	Model checkpoint (best validation loss)
hparams.yaml	Hyperparameters used for training
train.csv	Training metrics per epoch
val.csv	Validation metrics per epoch
vocab_stanza_hybrid.pkl	Vocabulary (for tree-based models)

Installation

# Clone ToxicThesis for full model implementations
git clone https://github.com/simo-corbo/ToxicThesis
cd ToxicThesis
pip install -r requirements.txt

# Or install dependencies directly
pip install torch transformers huggingface_hub fasttext-wheel stanza

Citation

@software{toxicthesis2025,
  title={ToxicThesis},
  author={Corbo, Simone},
  year={2025},
  url={https://github.com/simo-corbo/ToxicThesis}
}