LSTM - GEMINI-3.5-FLASH - Classification (3 classes)
Toxicity prediction model trained on the GEMINI-3.5-FLASH dataset.
| Property |
Value |
| Model |
LSTM |
| Task |
Classification (3 classes) |
| Dataset |
gemini-3.5-flash |
| Framework |
PyTorch / PyTorch Lightning |
Class: LSTMModel
from src.models.lstm import LSTMModel
model = LSTMModel(
input_dim: int = 300,
hidden_dim: int = 128,
num_layers: int = 2,
dropout: float = 0.3,
bidirectional: bool = True,
num_classes: int = 2,
loss_fn: str = 'auto',
lr: float = 0.001,
gradient_clip_norm: float = 1.0
)
Methods
| Method |
Description |
forward(inputs) |
Input: (batch, seq_len, input_dim). Returns raw logits. |
compute_loss(logits, targets) |
Computes loss based on num_classes. |
bin_targets(targets) |
Converts continuous [0,1] targets to class indices. |
from_config(config) |
Create model from config dict. |
load_from_checkpoint(path) |
Load model from checkpoint file. |
Usage with ToxicThesis (Recommended)
from huggingface_hub import hf_hub_download
import torch
import numpy as np
checkpoint_path = hf_hub_download(
repo_id="simocorbo/toxicthesis-gemini-3.5-flash-lstm-classification-3",
filename="checkpoints/best.pt"
)
from src.models.lstm import LSTMModel
model = LSTMModel.load_from_checkpoint(checkpoint_path, map_location='cpu')
model.eval()
from src.utils.fasttext_utils import load_fasttext_model
ft = load_fasttext_model('cc.en.300.bin')
def predict(text: str, max_len: int = 128) -> dict:
tokens = text.lower().split()[:max_len]
embeddings = [ft.get_word_vector(w) for w in tokens] or [np.zeros(300)]
while len(embeddings) < max_len:
embeddings.append(np.zeros(300))
x = torch.tensor(np.array(embeddings[:max_len]), dtype=torch.float32).unsqueeze(0)
with torch.no_grad():
logits = model(x)
if model.num_classes == 1:
score = torch.sigmoid(logits).item()
return {'score': score}
elif model.num_classes == 2:
prob = torch.sigmoid(logits).item()
return {'probability': prob, 'class': int(prob >= 0.5)}
else:
probs = torch.softmax(logits, dim=-1).squeeze().tolist()
return {'probabilities': probs, 'class': int(np.argmax(probs))}
result = predict("Your text here")
print(result)
Score Interpretation
| Output |
Range |
Meaning |
probabilities |
List[float] |
Probability distribution over 3 classes. |
class |
0 to 2 |
Predicted class (argmax of probabilities). |
Classes: 3 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
git clone https://github.com/simo-corbo/ToxicThesis
cd ToxicThesis
pip install -r requirements.txt
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}
}