"""Lightweight baseline text classifier for Phase 3.""" from __future__ import annotations import json import importlib import math import os import pickle import re import sys import tempfile from collections import Counter from pathlib import Path from typing import Any, Dict, Iterable, List, Protocol, Sequence, Tuple from prert.phase3.types import ClauseExample TOKEN_PATTERN = re.compile(r"[a-z0-9]{2,}") TFIDF_TOKEN_PATTERN = r"(?u)\b[a-zA-Z0-9][a-zA-Z0-9']+\b" DEFAULT_PRIVACYBERT_MODEL_NAME = "mukund/privbert" class TextClassifier(Protocol): def fit(self, examples: Iterable[ClauseExample]) -> None: ... def predict(self, text: str) -> str: ... def predict_proba(self, text: str) -> Dict[str, float]: ... def save(self, path: Path) -> None: ... class NaiveBayesTextClassifier: def __init__(self, labels: Sequence[str], alpha: float = 1.0) -> None: self.labels = list(labels) self.alpha = alpha self.class_doc_counts: Dict[str, int] = {label: 0 for label in self.labels} self.class_token_totals: Dict[str, int] = {label: 0 for label in self.labels} self.class_token_counts: Dict[str, Counter[str]] = {label: Counter() for label in self.labels} self.vocabulary: set[str] = set() self.total_docs = 0 def fit(self, examples: Iterable[ClauseExample]) -> None: for example in examples: if example.label not in self.class_doc_counts: continue tokens = tokenize(example.text) if not tokens: continue self.total_docs += 1 self.class_doc_counts[example.label] += 1 self.class_token_counts[example.label].update(tokens) self.class_token_totals[example.label] += len(tokens) self.vocabulary.update(tokens) def predict(self, text: str) -> str: scores = self._class_log_scores(text) return max(scores.items(), key=lambda item: item[1])[0] def predict_proba(self, text: str) -> Dict[str, float]: log_scores = self._class_log_scores(text) max_log = max(log_scores.values()) shifted = {label: math.exp(score - max_log) for label, score in log_scores.items()} total = sum(shifted.values()) if total <= 0: uniform = 1.0 / max(len(self.labels), 1) return {label: uniform for label in self.labels} return {label: shifted[label] / total for label in self.labels} def save(self, path: Path) -> None: payload = { "labels": self.labels, "alpha": self.alpha, "class_doc_counts": self.class_doc_counts, "class_token_totals": self.class_token_totals, "class_token_counts": {label: dict(counter) for label, counter in self.class_token_counts.items()}, "vocabulary": sorted(self.vocabulary), "total_docs": self.total_docs, } path.parent.mkdir(parents=True, exist_ok=True) with path.open("w", encoding="utf-8") as handle: json.dump(payload, handle, indent=2, ensure_ascii=False) handle.write("\n") @classmethod def load(cls, path: Path) -> "NaiveBayesTextClassifier": with path.open("r", encoding="utf-8") as handle: payload = json.load(handle) model = cls(labels=payload["labels"], alpha=float(payload["alpha"])) model.class_doc_counts = {k: int(v) for k, v in payload["class_doc_counts"].items()} model.class_token_totals = {k: int(v) for k, v in payload["class_token_totals"].items()} model.class_token_counts = { label: Counter({token: int(count) for token, count in counts.items()}) for label, counts in payload["class_token_counts"].items() } model.vocabulary = set(payload.get("vocabulary", [])) model.total_docs = int(payload.get("total_docs", sum(model.class_doc_counts.values()))) return model def _class_log_scores(self, text: str) -> Dict[str, float]: tokens = tokenize(text) vocab_size = max(len(self.vocabulary), 1) denominator_by_class = { label: self.class_token_totals[label] + (self.alpha * vocab_size) for label in self.labels } log_scores: Dict[str, float] = {} for label in self.labels: prior_numerator = self.class_doc_counts[label] + self.alpha prior_denominator = self.total_docs + (self.alpha * len(self.labels)) score = math.log(prior_numerator / max(prior_denominator, 1e-9)) token_counter = self.class_token_counts[label] denominator = denominator_by_class[label] for token in tokens: numerator = token_counter[token] + self.alpha score += math.log(numerator / max(denominator, 1e-9)) log_scores[label] = score return log_scores class TfidfLogisticRegressionClassifier: """TF-IDF + class-weighted logistic regression model. This model is designed to improve minority-class discrimination compared to the Naive Bayes baseline while preserving deterministic behavior. """ def __init__( self, labels: Sequence[str], random_state: int = 42, max_features: int = 20000, ngram_max: int = 2, min_df: int = 2, max_df: float = 0.95, c: float = 1.0, max_iter: int = 1000, class_weight: str | Dict[str, float] = "balanced", ) -> None: try: from sklearn.feature_extraction.text import TfidfVectorizer from sklearn.linear_model import LogisticRegression except ModuleNotFoundError as exc: raise RuntimeError( "Model type 'logreg_tfidf' requires scikit-learn. " "Install dependencies and re-run." ) from exc self.labels = list(labels) self.random_state = random_state self.max_features = max_features self.ngram_max = ngram_max self.min_df = min_df self.max_df = max_df self.c = c self.max_iter = max_iter self.class_weight = class_weight self.vectorizer = TfidfVectorizer( lowercase=True, token_pattern=TFIDF_TOKEN_PATTERN, max_features=max_features, ngram_range=(1, ngram_max), min_df=min_df, max_df=max_df, sublinear_tf=True, ) self.model = LogisticRegression( class_weight=class_weight, max_iter=max_iter, random_state=random_state, C=c, solver="lbfgs", ) self.is_fit = False def fit(self, examples: Iterable[ClauseExample]) -> None: texts: List[str] = [] labels: List[str] = [] for example in examples: if example.label not in self.labels: continue text = (example.text or "").strip() if not text: continue texts.append(text) labels.append(example.label) if not texts: raise ValueError("No training examples available for logistic regression model") x_train = self.vectorizer.fit_transform(texts) self.model.fit(x_train, labels) self.is_fit = True def predict(self, text: str) -> str: if not self.is_fit: raise RuntimeError("Model is not fit") features = self.vectorizer.transform([text]) prediction = self.model.predict(features)[0] return str(prediction) def predict_proba(self, text: str) -> Dict[str, float]: if not self.is_fit: raise RuntimeError("Model is not fit") features = self.vectorizer.transform([text]) proba = self.model.predict_proba(features)[0] probabilities = {label: 0.0 for label in self.labels} for label, value in zip(self.model.classes_, proba): probabilities[str(label)] = float(value) total = sum(probabilities.values()) if total <= 0: uniform = 1.0 / max(len(self.labels), 1) return {label: uniform for label in self.labels} return {label: probabilities[label] / total for label in self.labels} def save(self, path: Path) -> None: payload: Dict[str, Any] = { "labels": self.labels, "random_state": self.random_state, "max_features": self.max_features, "ngram_max": self.ngram_max, "min_df": self.min_df, "max_df": self.max_df, "c": self.c, "max_iter": self.max_iter, "class_weight": self.class_weight, "vectorizer": self.vectorizer, "model": self.model, } path.parent.mkdir(parents=True, exist_ok=True) with path.open("wb") as handle: pickle.dump(payload, handle) class _PrivacyBertTrainingDataset: def __init__(self, encodings: Dict[str, Sequence[Any]], labels: Sequence[int]) -> None: self.encodings = encodings self.labels = list(labels) def __len__(self) -> int: return len(self.labels) def __getitem__(self, index: int) -> Dict[str, Any]: item = {name: values[index] for name, values in self.encodings.items()} item["labels"] = self.labels[index] return item def _patch_multiprocess_resource_tracker() -> None: if os.name != "nt" or sys.version_info < (3, 12): return try: resource_tracker_module = importlib.import_module("multiprocess.resource_tracker") except ModuleNotFoundError: return resource_tracker_cls = getattr(resource_tracker_module, "ResourceTracker", None) if resource_tracker_cls is None or getattr(resource_tracker_cls, "_prert_safe_stop_locked", False): return def _safe_stop_locked( self, close=os.close, waitpid=os.waitpid, waitstatus_to_exitcode=os.waitstatus_to_exitcode, ): recursion_count = getattr(self._lock, "_recursion_count", None) recursion_depth = recursion_count() if callable(recursion_count) else None if callable(recursion_count) and recursion_depth > 1: return self._reentrant_call_error() if self._fd is None: return if self._pid is None: return close(self._fd) self._fd = None waitpid(self._pid, 0) self._pid = None resource_tracker_cls._stop_locked = _safe_stop_locked resource_tracker_cls._prert_safe_stop_locked = True class PrivacyBertClassifier: """PrivacyBERT-style transformer classifier backend. This backend is optional and only activated when model_type=privacybert. It requires transformers and torch to be installed. """ def __init__( self, labels: Sequence[str], model_name: str = DEFAULT_PRIVACYBERT_MODEL_NAME, random_state: int = 42, num_train_epochs: float = 2.0, batch_size: int = 8, learning_rate: float = 5e-5, max_length: int = 256, loss_type: str = "focal", focal_gamma: float = 2.0, label_smoothing_factor: float = 0.05, weight_decay: float = 0.01, warmup_steps: float = 0.1, early_stopping_patience: int = 1, ) -> None: try: torch = importlib.import_module("torch") transformers_module = importlib.import_module("transformers") except ModuleNotFoundError as exc: raise RuntimeError( "Model type 'privacybert' requires torch and transformers. " "Install dependencies and re-run." ) from exc _patch_multiprocess_resource_tracker() normalized_loss = (loss_type or "focal").strip().lower() if normalized_loss not in {"ce", "weighted_ce", "focal"}: raise ValueError( f"Unsupported loss_type '{loss_type}'. Expected one of: ce, weighted_ce, focal." ) self.labels = list(labels) self.model_name = model_name self.random_state = random_state self.num_train_epochs = num_train_epochs self.batch_size = batch_size self.learning_rate = learning_rate self.max_length = max_length self.loss_type = normalized_loss self.focal_gamma = float(focal_gamma) self.label_smoothing_factor = float(label_smoothing_factor) self.weight_decay = float(weight_decay) self.warmup_steps = float(warmup_steps) self.early_stopping_patience = int(early_stopping_patience) self._torch = torch self._training_dataset_cls = _PrivacyBertTrainingDataset self._trainer_cls = transformers_module.Trainer self._training_args_cls = transformers_module.TrainingArguments self._auto_tokenizer_cls = transformers_module.AutoTokenizer self._auto_model_cls = transformers_module.AutoModelForSequenceClassification self._early_stopping_cls = getattr(transformers_module, "EarlyStoppingCallback", None) self.label_to_id = {label: idx for idx, label in enumerate(self.labels)} self.id_to_label = {idx: label for label, idx in self.label_to_id.items()} self.tokenizer = self._auto_tokenizer_cls.from_pretrained(model_name) self.model = self._auto_model_cls.from_pretrained( model_name, num_labels=len(self.labels), label2id=self.label_to_id, id2label=self.id_to_label, ) self.is_fit = False def _has_accelerator(self) -> bool: accelerator_module = getattr(self._torch, "accelerator", None) if accelerator_module is not None and callable(getattr(accelerator_module, "is_available", None)): return bool(accelerator_module.is_available()) cuda_module = getattr(self._torch, "cuda", None) if cuda_module is not None and callable(getattr(cuda_module, "is_available", None)): return bool(cuda_module.is_available()) xpu_module = getattr(self._torch, "xpu", None) if xpu_module is not None and callable(getattr(xpu_module, "is_available", None)): return bool(xpu_module.is_available()) backends = getattr(self._torch, "backends", None) mps_module = getattr(backends, "mps", None) if backends is not None else None if mps_module is not None and callable(getattr(mps_module, "is_available", None)): return bool(mps_module.is_available()) return False def fit( self, examples: Iterable[ClauseExample], validation_examples: Iterable[ClauseExample] | None = None, ) -> None: texts: List[str] = [] labels: List[int] = [] for example in examples: label = str(example.label).strip().lower() if label not in self.label_to_id: continue text = (example.text or "").strip() if not text: continue texts.append(text) labels.append(self.label_to_id[label]) if not texts: raise ValueError("No training examples available for privacybert model") encodings = self.tokenizer( texts, truncation=True, padding="max_length", max_length=self.max_length, ) dataset = self._training_dataset_cls(encodings=encodings, labels=labels) eval_dataset = None if validation_examples is not None: val_texts: List[str] = [] val_labels: List[int] = [] for example in validation_examples: label = str(example.label).strip().lower() if label not in self.label_to_id: continue text = (example.text or "").strip() if not text: continue val_texts.append(text) val_labels.append(self.label_to_id[label]) if val_texts: val_encodings = self.tokenizer( val_texts, truncation=True, padding="max_length", max_length=self.max_length, ) eval_dataset = self._training_dataset_cls( encodings=val_encodings, labels=val_labels ) # Class weights for weighted_ce / focal: balanced inverse-frequency # over the actual training labels (min count guard avoids div-by-zero # when a class is absent from the train split). class_counts = [0] * len(self.labels) for idx in labels: class_counts[idx] += 1 total = sum(class_counts) n_classes = len(self.labels) class_weights = [ total / (n_classes * max(count, 1)) for count in class_counts ] class_weights_tensor = self._torch.tensor(class_weights, dtype=self._torch.float32) loss_type = self.loss_type focal_gamma = self.focal_gamma torch_mod = self._torch nn_module = getattr(torch_mod, "nn", None) functional_module = getattr(nn_module, "functional", None) if nn_module is not None else None if loss_type in {"weighted_ce", "focal"} and (nn_module is None or functional_module is None): raise RuntimeError( "torch.nn / torch.nn.functional are required for loss_type in {weighted_ce, focal}" ) class _Phase3WeightedTrainer(self._trainer_cls): # type: ignore[misc] def compute_loss(self, model, inputs, return_outputs=False, **_kwargs): # type: ignore[no-untyped-def] target_labels = inputs.get("labels") outputs = model(**inputs) logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] weights = class_weights_tensor.to(logits.device) if loss_type == "focal": log_probs = functional_module.log_softmax(logits, dim=-1) # type: ignore[union-attr] probs = log_probs.exp() target_idx = target_labels.long().unsqueeze(1) true_log_probs = log_probs.gather(1, target_idx).squeeze(1) true_probs = probs.gather(1, target_idx).squeeze(1) alpha_t = weights[target_labels.long()] focal_term = (1.0 - true_probs).pow(focal_gamma) loss = -(alpha_t * focal_term * true_log_probs).mean() else: # weighted_ce path (plain "ce" is handled via label_smoothing # in TrainingArguments and never reaches this subclass — we # only install the subclass when weights are needed). loss_fct = nn_module.CrossEntropyLoss(weight=weights) # type: ignore[union-attr] loss = loss_fct( logits.view(-1, logits.size(-1)), target_labels.view(-1), ) return (loss, outputs) if return_outputs else loss def _compute_metrics(eval_pred): # type: ignore[no-untyped-def] from sklearn.metrics import f1_score # local import: only used during real training predictions, label_ids = eval_pred # predictions may be tuple in some configurations if isinstance(predictions, tuple): predictions = predictions[0] preds = predictions.argmax(axis=-1) macro_f1 = f1_score(label_ids, preds, average="macro", zero_division=0) return {"macro_f1": float(macro_f1)} with tempfile.TemporaryDirectory(prefix="phase3-privacybert-") as tmpdir: training_args_kwargs: Dict[str, Any] = { "output_dir": tmpdir, "learning_rate": self.learning_rate, "per_device_train_batch_size": self.batch_size, "num_train_epochs": self.num_train_epochs, "dataloader_pin_memory": self._has_accelerator(), "logging_strategy": "no", "report_to": [], "seed": self.random_state, "data_seed": self.random_state, "weight_decay": self.weight_decay, "warmup_steps": self.warmup_steps, } # Label smoothing only applies to plain CE (loss is computed by HF # Trainer in that path). Skip under weighted_ce/focal since the # subclass owns the loss. if self.loss_type == "ce" and self.label_smoothing_factor > 0.0: training_args_kwargs["label_smoothing_factor"] = self.label_smoothing_factor if eval_dataset is not None: training_args_kwargs.update( { "eval_strategy": "epoch", "save_strategy": "epoch", "save_total_limit": 1, "per_device_eval_batch_size": self.batch_size, "load_best_model_at_end": True, "metric_for_best_model": "macro_f1", "greater_is_better": True, } ) else: training_args_kwargs["save_strategy"] = "no" training_args = self._training_args_cls(**training_args_kwargs) trainer_kwargs: Dict[str, Any] = { "model": self.model, "args": training_args, "train_dataset": dataset, } if eval_dataset is not None: trainer_kwargs["eval_dataset"] = eval_dataset trainer_kwargs["compute_metrics"] = _compute_metrics if self._early_stopping_cls is not None and self.early_stopping_patience > 0: trainer_kwargs["callbacks"] = [ self._early_stopping_cls( early_stopping_patience=self.early_stopping_patience ) ] # Plain CE goes through HF's default Trainer (its built-in loss # already supports label_smoothing_factor); weighted_ce/focal need # the custom subclass that owns compute_loss. trainer_class = ( self._trainer_cls if self.loss_type == "ce" else _Phase3WeightedTrainer ) trainer = trainer_class(**trainer_kwargs) trainer.train() self.model = trainer.model self.model.eval() self.is_fit = True def predict(self, text: str) -> str: probabilities = self.predict_proba(text) return max(probabilities.items(), key=lambda item: item[1])[0] def predict_proba(self, text: str) -> Dict[str, float]: if not self.is_fit: raise RuntimeError("Model is not fit") encoded = self.tokenizer( text, truncation=True, padding="max_length", max_length=self.max_length, return_tensors="pt", ) with self._torch.no_grad(): logits = self.model(**encoded).logits[0] probs = self._torch.softmax(logits, dim=-1).cpu().tolist() output = {label: 0.0 for label in self.labels} for idx, value in enumerate(probs): label = self.id_to_label.get(idx) if label is None: continue output[label] = float(value) total = sum(output.values()) if total <= 0: uniform = 1.0 / max(len(self.labels), 1) return {label: uniform for label in self.labels} return {label: output[label] / total for label in self.labels} def save(self, path: Path) -> None: save_dir = path if path.suffix == "" else (path.parent / path.stem) save_dir.mkdir(parents=True, exist_ok=True) self.model.save_pretrained(save_dir) self.tokenizer.save_pretrained(save_dir) metadata = { "labels": self.labels, "model_name": self.model_name, "random_state": self.random_state, "num_train_epochs": self.num_train_epochs, "batch_size": self.batch_size, "learning_rate": self.learning_rate, "max_length": self.max_length, "loss_type": self.loss_type, "focal_gamma": self.focal_gamma, "label_smoothing_factor": self.label_smoothing_factor, "weight_decay": self.weight_decay, "warmup_steps": self.warmup_steps, "early_stopping_patience": self.early_stopping_patience, } with (save_dir / "training_metadata.json").open("w", encoding="utf-8") as handle: json.dump(metadata, handle, indent=2, ensure_ascii=False) handle.write("\n") def train_classifier( examples: Sequence[ClauseExample], labels: Sequence[str], output_path: Path, model_type: str = "naive_bayes", random_state: int = 42, max_features: int = 20000, ngram_max: int = 2, min_df: int = 2, max_df: float = 0.95, c: float = 1.0, max_iter: int = 1000, privacybert_model_name: str = DEFAULT_PRIVACYBERT_MODEL_NAME, privacybert_epochs: float = 2.0, privacybert_batch_size: int = 8, privacybert_learning_rate: float = 5e-5, privacybert_max_length: int = 256, privacybert_loss_type: str = "focal", privacybert_focal_gamma: float = 2.0, privacybert_label_smoothing: float = 0.05, privacybert_weight_decay: float = 0.01, privacybert_warmup_steps: float = 0, privacybert_early_stopping_patience: int = 1, validation_examples: Sequence[ClauseExample] | None = None, ) -> Tuple[TextClassifier, Dict[str, float]]: selected_model = model_type.strip().lower() if selected_model in {"naive_bayes", "nb"}: model: TextClassifier = NaiveBayesTextClassifier(labels=labels) model_name = "multinomial_naive_bayes" elif selected_model in {"logreg_tfidf", "logistic_regression", "lr_tfidf"}: model = TfidfLogisticRegressionClassifier( labels=labels, random_state=random_state, max_features=max_features, ngram_max=ngram_max, min_df=min_df, max_df=max_df, c=c, max_iter=max_iter, ) model_name = "logreg_tfidf" elif selected_model in {"privacybert", "privacy_bert", "bert_privacy"}: model = PrivacyBertClassifier( labels=labels, model_name=privacybert_model_name, random_state=random_state, num_train_epochs=privacybert_epochs, batch_size=privacybert_batch_size, learning_rate=privacybert_learning_rate, max_length=privacybert_max_length, loss_type=privacybert_loss_type, focal_gamma=privacybert_focal_gamma, label_smoothing_factor=privacybert_label_smoothing, weight_decay=privacybert_weight_decay, warmup_steps=privacybert_warmup_steps, early_stopping_patience=privacybert_early_stopping_patience, ) model_name = "privacybert" else: raise ValueError(f"Unsupported model_type '{model_type}'") if isinstance(model, PrivacyBertClassifier): model.fit(examples, validation_examples=validation_examples) else: model.fit(examples) model.save(output_path) vocabulary_size = 0.0 if isinstance(model, NaiveBayesTextClassifier): vocabulary_size = float(len(model.vocabulary)) elif isinstance(model, TfidfLogisticRegressionClassifier): vocabulary_size = float(len(model.vectorizer.vocabulary_)) summary = { "model_type": model_name, "training_examples": float(len(examples)), "vocabulary_size": vocabulary_size, "labels": float(len(labels)), "backbone_model_name": getattr(model, "model_name", ""), } return model, summary def tokenize(text: str) -> List[str]: lowered = text.lower() return TOKEN_PATTERN.findall(lowered)