Toxicity prediction model trained on the GEMINI-3.5-FLASH dataset.
1from src.models.linear import LinearModel
2
3model = LinearModel(
4 input_dim: int, # Dimension of input embeddings (default: 300 for FastText)
5 num_classes: int = 2, # 1=regression, 2=binary, 3+=multi-class
6 loss_fn: str = 'auto', # 'auto', 'mse', 'bce', 'cross_entropy'
7 dropout: float = 0.0, # Dropout rate
8 l2_lambda: float = 0.01, # L2 regularization strength
9 learning_rate: float = 0.01,
10 max_epochs: int = 100,
11 batch_size: int = 512,
12 device: str = None # 'cuda', 'mps', 'cpu', or None for auto
13)
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-linear-classification-2",
11 filename="checkpoints/best.pt"
12)
13
14# 3. Import and load model from ToxicThesis
15from src.models.linear import LinearModel
16
17# Load checkpoint to get hyperparameters
18checkpoint = torch.load(checkpoint_path, map_location='cpu', weights_only=False)
19hparams = checkpoint.get('hyper_parameters', {})
20
21# Create model with same config
22model = LinearModel(
23 input_dim=hparams.get('input_dim', 300),
24 num_classes=hparams.get('num_classes', 2),
25 dropout=hparams.get('dropout', 0.0),
26 l2_lambda=hparams.get('l2_lambda', 0.01)
27)
28
29# Load trained weights
30state_dict = checkpoint.get('state_dict', checkpoint.get('model_state_dict', checkpoint))
31model.model.load_state_dict(state_dict, strict=False)
32model.model.eval()
33
34# 4. Get predictions using built-in methods
35import numpy as np
36from src.utils.fasttext_utils import load_fasttext_model
37
38ft = load_fasttext_model('cc.en.300.bin')
39
40def get_embedding(text: str) -> np.ndarray:
41 tokens = text.lower().split()
42 embeddings = [ft.get_word_vector(w) for w in tokens]
43 return np.mean(embeddings, axis=0) if embeddings else np.zeros(300)
44
45# Single prediction
46text = "Your text here"
47emb = get_embedding(text)
48X = torch.tensor(emb, dtype=torch.float32).unsqueeze(0).to(model.device)
49
50with torch.no_grad():
51 # Use predict_proba for probabilities
52 probs = model.model.predict_proba(X)
53 print(f"Probabilities: {probs}")
54
55 # Or use forward for raw logits
56 logits = model.model(X)
57 print(f"Logits: {logits}")
1from huggingface_hub import hf_hub_download
2import torch
3import torch.nn as nn
4import fasttext
5import numpy as np
6
7# 1. Download checkpoint
8checkpoint_path = hf_hub_download(
9 repo_id="simocorbo/toxicthesis-gemini-3.5-flash-linear-classification-2",
10 filename="checkpoints/best.pt"
11)
12
13# 2. Load FastText embeddings
14ft = fasttext.load_model('cc.en.300.bin')
15
16# 3. Define minimal model class
17class LinearClassifier(nn.Module):
18 def __init__(self, input_dim: int, num_classes: int):
19 super().__init__()
20 self.num_classes = num_classes
21 output_dim = 1 if num_classes <= 2 else num_classes
22 self.output = nn.Linear(input_dim, output_dim)
23
24 def forward(self, x):
25 return self.output(x)
26
27# 4. Load checkpoint
28checkpoint = torch.load(checkpoint_path, map_location='cpu', weights_only=False)
29hparams = checkpoint.get('hyper_parameters', {})
30num_classes = hparams.get('num_classes', 2)
31
32model = LinearClassifier(input_dim=300, num_classes=num_classes)
33state_dict = checkpoint.get('state_dict', checkpoint.get('model_state_dict', checkpoint))
34state_dict = {k.replace('model.', '').replace('linear.', ''): v for k, v in state_dict.items()}
35model.load_state_dict(state_dict, strict=False)
36model.eval()
37
38# 5. Inference function
39def predict(text: str) -> dict:
40 tokens = text.lower().split()
41 emb = np.mean([ft.get_word_vector(w) for w in tokens], axis=0) if tokens else np.zeros(300)
42 x = torch.tensor(emb, dtype=torch.float32).unsqueeze(0)
43
44 with torch.no_grad():
45 logits = model(x)
46 if num_classes == 1:
47 score = torch.sigmoid(logits).item()
48 return {'score': score}
49 elif num_classes == 2:
50 prob = torch.sigmoid(logits).item()
51 return {'probability': prob, 'class': int(prob >= 0.5)}
52 else:
53 probs = torch.softmax(logits, dim=-1).squeeze().tolist()
54 return {'probabilities': probs, 'class': int(np.argmax(probs))}
55
56result = predict("Your text here")
57print(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}
