Toxicity prediction model trained on the GEMINI-3.5-FLASH dataset.
1from src.models.roberta import RobertaModel
2
3model = RobertaModel(
4 model_name: str = 'roberta-base', # HuggingFace model name
5 num_classes: int = 2, # 1=regression, 2=binary, 3+=multi-class
6 loss_fn: str = 'auto',
7 freeze_layers: int = 0, # Number of layers to freeze
8 dropout: float = 0.1,
9 lr: float = 2e-5,
10 gradient_checkpointing: bool = True
11)
1# 1. Clone ToxicThesis repository
2# git clone https://github.com/simo-corbo/ToxicThesis
3# cd ToxicThesis && pip install -r requirements.txt
4
5from huggingface_hub import hf_hub_download
6import torch
7
8# 2. Download checkpoint
9checkpoint_path = hf_hub_download(
10 repo_id="simocorbo/toxicthesis-gemini-3.5-flash-roberta-classification-2",
11 filename="checkpoints/best.pt"
12)
13
14# 3. Import and load model from ToxicThesis
15from src.models.roberta import RobertaModel
16
17checkpoint = torch.load(checkpoint_path, map_location='cpu', weights_only=False)
18hparams = checkpoint.get('hyper_parameters', {})
19
20model = RobertaModel(
21 model_name=hparams.get('model_name', 'roberta-base'),
22 num_classes=hparams.get('num_classes', 2),
23 dropout=hparams.get('dropout', 0.1)
24)
25model.load_state_dict(checkpoint.get('state_dict', checkpoint), strict=False)
26model.eval()
27
28# 4. Use the built-in predict_text method (easiest!)
29result = model.predict_text("Your text here")
30print(result)
31# Output: {'score': 0.7234, 'prediction': 'Toxicity score: 0.7234', ...}
32
33# 5. Multiple predictions
34texts = ["Hello friend", "You are terrible", "Have a nice day"]
35for text in texts:
36 result = model.predict_text(text)
37 print(f"{text}: {result['score']:.4f}")
1from huggingface_hub import hf_hub_download
2import torch
3import torch.nn as nn
4from transformers import AutoModel, AutoTokenizer
5
6# 1. Download checkpoint
7checkpoint_path = hf_hub_download(
8 repo_id="simocorbo/toxicthesis-gemini-3.5-flash-roberta-classification-2",
9 filename="checkpoints/best.pt"
10)
11
12# 2. Load checkpoint
13checkpoint = torch.load(checkpoint_path, map_location='cpu', weights_only=False)
14hparams = checkpoint.get('hyper_parameters', {})
15model_name = hparams.get('model_name', 'roberta-base')
16num_classes = hparams.get('num_classes', 2)
17
18# 3. Define model class
19class RobertaClassifier(nn.Module):
20 def __init__(self, model_name, num_classes, dropout=0.1):
21 super().__init__()
22 self.num_classes = num_classes
23 self.roberta = AutoModel.from_pretrained(model_name, add_pooling_layer=False)
24 hidden_size = self.roberta.config.hidden_size
25 self.dropout = nn.Dropout(dropout)
26 self.classifier = nn.Linear(hidden_size, 1 if num_classes <= 2 else num_classes)
27
28 def forward(self, input_ids, attention_mask=None):
29 outputs = self.roberta(input_ids=input_ids, attention_mask=attention_mask)
30 pooled = self.dropout(outputs.last_hidden_state[:, 0, :])
31 return self.classifier(pooled)
32
33# 4. Load model
34model = RobertaClassifier(model_name, num_classes)
35state_dict = checkpoint.get('state_dict', checkpoint)
36model.load_state_dict(state_dict, strict=False)
37model.eval()
38
39# 5. Inference
40tokenizer = AutoTokenizer.from_pretrained(model_name)
41
42def predict(text: str) -> dict:
43 inputs = tokenizer(text, return_tensors='pt', truncation=True, max_length=512)
44 with torch.no_grad():
45 logits = model(inputs['input_ids'], inputs.get('attention_mask'))
46 if num_classes == 1:
47 return {'score': torch.sigmoid(logits).item()}
48 elif num_classes == 2:
49 p = torch.sigmoid(logits).item()
50 return {'probability': p, 'class': int(p >= 0.5)}
51 else:
52 probs = torch.softmax(logits, dim=-1).squeeze().tolist()
53 return {'probabilities': probs, 'class': int(torch.argmax(torch.tensor(probs)))}
54
55result = predict("Your text here")
56print(result)
1# Clone ToxicThesis for full model implementations
2git clone https://github.com/simo-corbo/ToxicThesis
3cd ToxicThesis
4pip install -r requirements.txt
5
6# Or install dependencies directly
7pip install torch transformers huggingface_hub fasttext-wheel stanza
1@software{toxicthesis2025,
2 title={ToxicThesis},
3 author={Corbo, Simone},
4 year={2025},
5 url={https://github.com/simo-corbo/ToxicThesis}
6}
