Deploy algospeak classifier — model loaded from HF Hub at runtime

Dockerfile

@@ -1,4 +1,4 @@
-FROM python:3.13.5-slim
+FROM python:3.12-slim
 
 WORKDIR /app
 
@@ -9,12 +9,12 @@ RUN apt-get update && apt-get install -y \
     && rm -rf /var/lib/apt/lists/*
 
 COPY requirements.txt ./
-COPY src/ ./src/
+RUN pip install --no-cache-dir -r requirements.txt
 
-RUN pip3 install -r requirements.txt
+COPY . .
 
 EXPOSE 8501
 
 HEALTHCHECK CMD curl --fail http://localhost:8501/_stcore/health
 
-ENTRYPOINT ["streamlit", "run", "src/streamlit_app.py", "--server.port=8501", "--server.address=0.0.0.0"]
+ENTRYPOINT ["streamlit", "run", "app.py", "--server.port=8501", "--server.address=0.0.0.0"]

app.py

@@ -0,0 +1,70 @@
+"""
+app.py — Algospeak Classifier demo
+
+Streamlit UI for the dual BERTweet model.
+Type a social media post and see the predicted class + confidence scores.
+"""
+
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).parent / "poc" / "src"))
+
+import yaml
+import torch
+import numpy as np
+import emoji
+import streamlit as st
+from transformers import AutoTokenizer
+from huggingface_hub import hf_hub_download
+
+from inference import load_unsupervised_encoder, classify_text
+
+BASE_DIR = Path(__file__).parent
+MODEL_REPO = "timagonch/algospeak-classifier-model"
+
+CLASS_COLORS = {
+    "Allowed":            "green",
+    "Offensive Language": "red",
+    "Mature Content":     "orange",
+    "Algospeak":          "violet",
+}
+
+
+@st.cache_resource(show_spinner="Loading model...")
+def load_model():
+    with open(BASE_DIR / "poc" / "config.yaml") as f:
+        cfg = yaml.safe_load(f)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    checkpoint_path = hf_hub_download(repo_id=MODEL_REPO, filename="best_model.pt")
+    prototypes_path = hf_hub_download(repo_id=MODEL_REPO, filename="prototypes.npy")
+
+    encoder = load_unsupervised_encoder(checkpoint_path, cfg, device)
+    prototypes = np.load(prototypes_path)
+    tokenizer = AutoTokenizer.from_pretrained(cfg["model_name"], use_fast=False)
+    return encoder, prototypes, tokenizer, cfg, device
+
+
+# ─────────────────────────────────────────────────────────────────────
+# UI
+# ─────────────────────────────────────────────────────────────────────
+
+st.title("Algospeak Classifier")
+st.caption("Dual BERTweet model · type a social media post to classify it.")
+
+text = st.text_area("Post text", height=120, placeholder="Type something here...")
+
+if st.button("Classify", type="primary") and text.strip():
+    encoder, prototypes, tokenizer, cfg, device = load_model()
+    result = classify_text(text, encoder, prototypes, tokenizer, cfg["max_length"], device)
+
+    label = result["predicted_label"]
+    color = CLASS_COLORS[label]
+
+    st.markdown(f"## :{color}[{label}]")
+    st.divider()
+
+    st.write("**Confidence scores:**")
+    for name, score in sorted(result["scores"].items(), key=lambda x: -x[1]):
+        st.progress(float(score), text=f"{name}: {score:.1%}")

poc/config.yaml

@@ -0,0 +1,35 @@
+# Dual BERTweet Configuration
+
+# Classes
+num_classes: 4
+class_labels:
+  0: "Allowed"
+  1: "Offensive Language"
+  2: "Mature Content"
+  3: "Algospeak"
+
+# Model
+model_name: "vinai/bertweet-base"
+embedding_dim: 768
+max_length: 128
+
+# Training
+batch_size: 32
+learning_rate: 2.0e-5
+weight_decay: 0.01
+num_epochs: 20
+warmup_steps: 200
+early_stopping_patience: 5
+fp16: true
+gradient_clip: 1.0
+
+# Loss
+temperature: 0.07
+
+# Paths (relative to project root)
+train_csv: "data/splits/train.csv"
+val_csv:   "data/splits/val.csv"
+test_csv:  "data/splits/test.csv"
+prepared_dir:   "poc/data/prepared"
+checkpoint_dir: "poc/checkpoints"
+results_dir:    "poc/results"

poc/src/inference.py

@@ -0,0 +1,328 @@
+"""
+inference.py
+
+Inference and full evaluation for the dual BERTweet model.
+
+Inference uses only the unsupervised encoder:
+  1. Build class prototypes from the training set (average embedding per class).
+  2. For a new post: encode -> cosine similarity to each prototype -> argmax = class.
+
+Evaluation produces:
+  - Accuracy (overall + per-class)
+  - Precision, Recall, F1 (per-class, macro, weighted)
+  - Confusion matrix (saved as PNG)
+  - ROC curves + AUC per class (saved as PNG)
+  - Full metrics saved to JSON
+
+Usage:
+    uv run python poc/src/inference.py
+"""
+
+import sys
+import json
+import yaml
+import torch
+import torch.nn.functional as F
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+import matplotlib
+import emoji
+matplotlib.use("Agg")  # non-interactive backend for saving figures
+
+from pathlib import Path
+from torch.utils.data import TensorDataset, DataLoader
+from transformers import AutoTokenizer
+from sklearn.metrics import (
+    accuracy_score,
+    classification_report,
+    confusion_matrix,
+    roc_curve,
+    auc,
+)
+
+sys.path.insert(0, str(Path(__file__).parent))
+from model import DualEncoderModel, BERTweetEncoder
+
+BASE_DIR = Path(__file__).resolve().parent.parent.parent
+
+CLASS_PREFIX = {
+    0: "Allowed:",
+    1: "Offensive Language:",
+    2: "Mature Content:",
+    3: "Algospeak:",
+}
+
+CLASS_NAMES = ["Allowed", "Offensive Language", "Mature Content", "Algospeak"]
+
+
+def load_config() -> dict:
+    with open(BASE_DIR / "poc" / "config.yaml") as f:
+        return yaml.safe_load(f)
+
+
+def load_unsupervised_encoder(ckpt_path: Path, cfg: dict, device: torch.device):
+    """Load the full dual model from checkpoint, return only the unsupervised encoder."""
+    model = DualEncoderModel(cfg["model_name"], cfg["temperature"])
+    ckpt  = torch.load(ckpt_path, map_location=device, weights_only=True)
+    model.load_state_dict(ckpt["model_state_dict"])
+    model = model.to(device)
+    model.eval()
+    print(f"Loaded checkpoint from epoch {ckpt['epoch']} (val_loss={ckpt['val_loss']:.4f})")
+    return model.unsupervised
+
+
+def load_dataset(path: Path) -> TensorDataset:
+    data = torch.load(path, map_location="cpu", weights_only=True)
+    return TensorDataset(
+        data["unsup_ids"],
+        data["unsup_mask"],
+        data["labels"],
+    )
+
+
+def get_embeddings(
+    encoder:  BERTweetEncoder,
+    dataset:  TensorDataset,
+    batch_sz: int,
+    device:   torch.device,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Run all samples through the unsupervised encoder. Returns (embeddings, labels)."""
+    loader = DataLoader(dataset, batch_size=batch_sz, shuffle=False, num_workers=2)
+    all_embs, all_labels = [], []
+
+    with torch.no_grad():
+        for unsup_ids, unsup_mask, labels in loader:
+            unsup_ids  = unsup_ids.to(device)
+            unsup_mask = unsup_mask.to(device)
+            embs = encoder(unsup_ids, unsup_mask)
+            all_embs.append(embs.cpu().numpy())
+            all_labels.append(labels.numpy())
+
+    return np.vstack(all_embs), np.concatenate(all_labels)
+
+
+def build_prototypes(
+    embeddings: np.ndarray,
+    labels:     np.ndarray,
+    num_classes: int,
+) -> np.ndarray:
+    """Average embedding per class -> [num_classes, D] prototype matrix."""
+    D = embeddings.shape[1]
+    prototypes = np.zeros((num_classes, D), dtype=np.float32)
+    for cls in range(num_classes):
+        mask = labels == cls
+        if mask.sum() > 0:
+            proto = embeddings[mask].mean(axis=0)
+            prototypes[cls] = proto / (np.linalg.norm(proto) + 1e-8)
+    return prototypes
+
+
+def predict(
+    embeddings: np.ndarray,
+    prototypes: np.ndarray,
+) -> tuple[np.ndarray, np.ndarray]:
+    """
+    Cosine similarity of each embedding to each prototype.
+    Returns (predicted_labels, score_matrix [N, num_classes]).
+    Scores are softmax-normalized cosine similarities — used for ROC curves.
+    """
+    # cosine similarity: embeddings are already L2-normalized, prototypes also normalized
+    sim = embeddings @ prototypes.T                           # [N, num_classes]
+    scores = torch.softmax(torch.tensor(sim / 0.1), dim=-1).numpy()  # [N, num_classes]
+    preds  = sim.argmax(axis=1)
+    return preds, scores
+
+
+# ─────────────────────────────────────────────────────────────────────
+# Plotting helpers
+# ─────────────────────────────────────────────────────────────────────
+
+def plot_confusion_matrix(y_true, y_pred, out_path: Path):
+    cm = confusion_matrix(y_true, y_pred)
+    fig, ax = plt.subplots(figsize=(7, 6))
+    im = ax.imshow(cm, interpolation="nearest", cmap=plt.cm.Blues)
+    plt.colorbar(im, ax=ax)
+
+    ax.set_xticks(range(len(CLASS_NAMES)))
+    ax.set_yticks(range(len(CLASS_NAMES)))
+    ax.set_xticklabels(CLASS_NAMES, rotation=30, ha="right", fontsize=9)
+    ax.set_yticklabels(CLASS_NAMES, fontsize=9)
+    ax.set_xlabel("Predicted")
+    ax.set_ylabel("True")
+    ax.set_title("Confusion Matrix")
+
+    thresh = cm.max() / 2.0
+    for i in range(cm.shape[0]):
+        for j in range(cm.shape[1]):
+            ax.text(j, i, str(cm[i, j]),
+                    ha="center", va="center",
+                    color="white" if cm[i, j] > thresh else "black", fontsize=10)
+
+    plt.tight_layout()
+    plt.savefig(out_path, dpi=150)
+    plt.close()
+    print(f"  Confusion matrix saved -> {out_path}")
+
+
+def plot_roc_curves(y_true, scores, num_classes: int, out_path: Path):
+    fig, ax = plt.subplots(figsize=(8, 6))
+    colors  = ["#e41a1c", "#377eb8", "#4daf4a", "#984ea3"]
+
+    for cls in range(num_classes):
+        y_bin = (y_true == cls).astype(int)
+        fpr, tpr, _ = roc_curve(y_bin, scores[:, cls])
+        roc_auc     = auc(fpr, tpr)
+        ax.plot(fpr, tpr, color=colors[cls], lw=2,
+                label=f"{CLASS_NAMES[cls]} (AUC={roc_auc:.3f})")
+
+    ax.plot([0, 1], [0, 1], "k--", lw=1)
+    ax.set_xlabel("False Positive Rate")
+    ax.set_ylabel("True Positive Rate")
+    ax.set_title("ROC Curves (One-vs-Rest)")
+    ax.legend(loc="lower right", fontsize=9)
+    plt.tight_layout()
+    plt.savefig(out_path, dpi=150)
+    plt.close()
+    print(f"  ROC curves saved -> {out_path}")
+
+
+# ─────────────────────────────────────────────────────────────────────
+# Main evaluation
+# ─────────────────────────────────────────────────────────────────────
+
+def evaluate_split(
+    encoder:    BERTweetEncoder,
+    prototypes: np.ndarray,
+    split:      str,
+    cfg:        dict,
+    device:     torch.device,
+    results_dir: Path,
+) -> dict:
+    print(f"\n--- Evaluating {split} split ---")
+    dataset = load_dataset(BASE_DIR / cfg["prepared_dir"] / f"{split}.pt")
+    embs, labels = get_embeddings(encoder, dataset, cfg["batch_size"], device)
+    preds, scores = predict(embs, prototypes)
+
+    # Save per-sample predictions CSV
+    csv_df = pd.read_csv(BASE_DIR / cfg[f"{split}_csv"])
+    csv_df = csv_df.dropna(subset=["text"]).reset_index(drop=True)
+    pred_df = pd.DataFrame({
+        "text":            csv_df["text"].astype(str),
+        "true_label":      [CLASS_NAMES[i] for i in labels],
+        "predicted_label": [CLASS_NAMES[i] for i in preds],
+        "correct":         labels == preds,
+    })
+    pred_df.to_csv(results_dir / f"predictions_{split}.csv", index=False)
+    print(f"  Predictions saved -> {results_dir / f'predictions_{split}.csv'}")
+
+    acc    = accuracy_score(labels, preds)
+    report = classification_report(
+        labels, preds, target_names=CLASS_NAMES, output_dict=True
+    )
+
+    print(f"  Accuracy: {acc:.4f}")
+    print(classification_report(labels, preds, target_names=CLASS_NAMES, digits=4))
+
+    plot_confusion_matrix(labels, preds, results_dir / f"confusion_matrix_{split}.png")
+    plot_roc_curves(labels, scores, cfg["num_classes"], results_dir / f"roc_curves_{split}.png")
+
+    aucs = {}
+    for cls in range(cfg["num_classes"]):
+        y_bin = (labels == cls).astype(int)
+        fpr, tpr, _ = roc_curve(y_bin, scores[:, cls])
+        aucs[CLASS_NAMES[cls]] = round(auc(fpr, tpr), 4)
+
+    return {
+        "split":           split,
+        "accuracy":        round(acc, 4),
+        "macro_f1":        round(report["macro avg"]["f1-score"], 4),
+        "weighted_f1":     round(report["weighted avg"]["f1-score"], 4),
+        "per_class":       {
+            CLASS_NAMES[i]: {
+                "precision": round(report[CLASS_NAMES[i]]["precision"], 4),
+                "recall":    round(report[CLASS_NAMES[i]]["recall"], 4),
+                "f1":        round(report[CLASS_NAMES[i]]["f1-score"], 4),
+            }
+            for i in range(cfg["num_classes"])
+        },
+        "auc_per_class":   aucs,
+        "mean_auc":        round(np.mean(list(aucs.values())), 4),
+    }
+
+
+def classify_text(text: str, encoder, prototypes, tokenizer, max_length, device) -> dict:
+    """Classify a single raw text string. Returns predicted class and similarity scores."""
+    enc = tokenizer(
+        emoji.demojize(text), padding="max_length", truncation=True,
+        max_length=max_length, return_tensors="pt",
+    )
+    with torch.no_grad():
+        emb = encoder(enc["input_ids"].to(device), enc["attention_mask"].to(device))
+    emb = emb.cpu().numpy()
+
+    sim    = emb @ prototypes.T
+    scores = torch.softmax(torch.tensor(sim / 0.1), dim=-1).numpy()[0]
+    pred   = int(sim.argmax())
+
+    return {
+        "predicted_class": pred,
+        "predicted_label": CLASS_NAMES[pred],
+        "scores":          {CLASS_NAMES[i]: round(float(scores[i]), 4)
+                            for i in range(len(CLASS_NAMES))},
+    }
+
+
+def main():
+    cfg    = load_config()
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Device: {device}")
+
+    ckpt_dir    = BASE_DIR / cfg["checkpoint_dir"]
+    results_dir = BASE_DIR / cfg["results_dir"]
+    results_dir.mkdir(parents=True, exist_ok=True)
+
+    # Load unsupervised encoder
+    encoder = load_unsupervised_encoder(ckpt_dir / "best_model.pt", cfg, device)
+
+    # Build prototypes from training set
+    print("\nBuilding class prototypes from training set...")
+    train_ds = load_dataset(BASE_DIR / cfg["prepared_dir"] / "train.pt")
+    train_embs, train_labels = get_embeddings(encoder, train_ds, cfg["batch_size"], device)
+    prototypes = build_prototypes(train_embs, train_labels, cfg["num_classes"])
+    np.save(results_dir / "prototypes.npy", prototypes)
+    print(f"  Prototypes saved -> {results_dir / 'prototypes.npy'}")
+
+    # Evaluate val and test splits
+    all_results = []
+    for split in ["val", "test"]:
+        result = evaluate_split(encoder, prototypes, split, cfg, device, results_dir)
+        all_results.append(result)
+
+    # Save metrics
+    metrics_path = results_dir / "metrics.json"
+    with open(metrics_path, "w") as f:
+        json.dump(all_results, f, indent=2)
+    print(f"\nAll metrics saved -> {metrics_path}")
+
+    # Summary
+    print("\n=== SUMMARY ===")
+    for r in all_results:
+        print(f"{r['split']:6s} | acc={r['accuracy']:.4f} | macro_f1={r['macro_f1']:.4f} | mean_auc={r['mean_auc']:.4f}")
+
+    # Quick example inference
+    print("\n=== Example inference ===")
+    tokenizer = AutoTokenizer.from_pretrained(cfg["model_name"], use_fast=False)
+    examples = [
+        "I had a great day today, went for a walk in the park.",
+        "I'm going to k!ll that n!gga if he shows up again.",
+        "she posted an onlyfans link in her bio",
+        "gonna unalive myself fr fr cant take this anymore",
+    ]
+    for text in examples:
+        result = classify_text(text, encoder, prototypes, tokenizer, cfg["max_length"], device)
+        print(f"  [{result['predicted_label']}] {text[:70]}")
+
+
+if __name__ == "__main__":
+    main()

poc/src/model.py

@@ -0,0 +1,116 @@
+"""
+model.py
+
+Dual BERTweet architecture for algospeak content moderation.
+
+Two independent BERTweet encoders trained jointly with supervised InfoNCE loss:
+  - supervised encoder:   receives "[CLASS_LABEL]: text" — class-aware during training
+  - unsupervised encoder: receives raw text only — the inference model
+
+At inference, only the unsupervised encoder is used. Its embeddings are compared
+to class prototypes (built from training data) via cosine similarity.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import AutoModel
+
+
+class BERTweetEncoder(nn.Module):
+    """
+    Wraps vinai/bertweet-base and returns an L2-normalized CLS token embedding.
+    """
+
+    def __init__(self, model_name: str):
+        super().__init__()
+        self.bert = AutoModel.from_pretrained(model_name)
+
+    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
+        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
+        cls_emb = outputs.last_hidden_state[:, 0, :]  # [B, 768]
+        return F.normalize(cls_emb, dim=-1)           # L2 normalize -> cosine-ready
+
+
+class DualEncoderModel(nn.Module):
+    """
+    Two independent BERTweet encoders trained with supervised InfoNCE loss.
+
+    supervised encoder:
+        Input: "[CLASS_LABEL]: <text>"  (e.g. "Offensive Language: I hate you")
+        Produces class-aware embeddings during training.
+        Discarded after training.
+
+    unsupervised encoder:
+        Input: raw text
+        Trained (via InfoNCE) to match the supervised encoder's embedding space.
+        Used exclusively at inference.
+    """
+
+    def __init__(self, model_name: str, temperature: float):
+        super().__init__()
+        self.supervised   = BERTweetEncoder(model_name)
+        self.unsupervised = BERTweetEncoder(model_name)
+        self.temperature  = temperature
+
+    def forward(
+        self,
+        sup_ids:    torch.Tensor,
+        sup_mask:   torch.Tensor,
+        unsup_ids:  torch.Tensor,
+        unsup_mask: torch.Tensor,
+        labels:     torch.Tensor,
+    ):
+        e_s = self.supervised(sup_ids, sup_mask)        # [B, D]
+        e_u = self.unsupervised(unsup_ids, unsup_mask)  # [B, D]
+        loss = supervised_infonce_loss(e_s, e_u, labels, self.temperature)
+        return loss, e_s, e_u
+
+
+def supervised_infonce_loss(
+    e_s:         torch.Tensor,
+    e_u:         torch.Tensor,
+    labels:      torch.Tensor,
+    temperature: float,
+) -> torch.Tensor:
+    """
+    Cross-encoder supervised InfoNCE loss.
+
+    For each unsupervised embedding e_u_i:
+      Positives: all supervised embeddings e_s_j where label_j == label_i
+      Negatives: all supervised embeddings e_s_j where label_j != label_i
+
+    Loss = mean_i [ -log( sum_{j: pos} exp(sim_ij/τ) / sum_j exp(sim_ij/τ) ) ]
+
+    Both e_s and e_u are L2-normalized so sim = dot product = cosine similarity.
+
+    Args:
+        e_s:         [B, D] supervised encoder embeddings
+        e_u:         [B, D] unsupervised encoder embeddings
+        labels:      [B]    integer class labels
+        temperature: scalar τ (typically 0.07)
+
+    Returns:
+        Scalar loss.
+    """
+    # Similarity matrix: unsupervised queries supervised keys — [B, B]
+    sim = torch.mm(e_u, e_s.T) / temperature
+
+    # Positive mask: True where label_j == label_i — [B, B]
+    pos_mask = (labels.unsqueeze(1) == labels.unsqueeze(0)).float()
+
+    # Numerical stability: subtract row max before exp
+    sim_max, _ = sim.max(dim=1, keepdim=True)
+    sim = sim - sim_max.detach()
+
+    exp_sim  = torch.exp(sim)
+    pos_sum  = (exp_sim * pos_mask).sum(dim=1)   # [B]
+    all_sum  = exp_sim.sum(dim=1)                # [B]
+
+    # Skip samples with no positives in this batch (shouldn't happen at batch_size >= num_classes)
+    valid = pos_sum > 0
+    if not valid.any():
+        return torch.tensor(0.0, requires_grad=True, device=e_s.device)
+
+    loss = -torch.log(pos_sum[valid] / all_sum[valid])
+    return loss.mean()

requirements.txt

@@ -1,3 +1,9 @@
-altair
-pandas
-streamlit
+torch
+transformers>=4.57.1
+streamlit>=1.56.0
+numpy
+pyyaml>=6.0.3
+emoji==0.6.0
+scikit-learn>=1.8.0
+sentencepiece
+huggingface_hub