# bert-class.py # Model training script for BERT-based sequence classification # This script uses the Hugging Face Transformers library to train a BERT model # for binary classification tasks, such as detecting potential violations in housing data. import os import numpy as np from dataclasses import dataclass from typing import Dict, Any, Optional import torch import torch.nn as nn from datasets import load_dataset from transformers import ( AutoTokenizer, AutoModelForSequenceClassification, DataCollatorWithPadding, TrainingArguments, Trainer, ) import evaluate from sklearn.utils.class_weight import compute_class_weight # ============================= # Settings # ============================= MODEL_NAME = os.getenv("MODEL_NAME", "google/bert_uncased_L-2_H-128_A-2") HUB_REPO = os.getenv("HUB_REPO", "tlogandesigns/fairhousing-bert-tiny") MAX_LEN = int(os.getenv("MAX_LEN", "256")) TRAIN_PATH = os.getenv("TRAIN_PATH", "train.csv") VAL_PATH = os.getenv("VAL_PATH", "val.csv") # Use binary labels 0/1 id2label = {0: "Compliant", 1: "Potential Violation"} label2id = {v: k for k, v in id2label.items()} # Metrics accuracy = evaluate.load("accuracy") precision = evaluate.load("precision") recall = evaluate.load("recall") f1 = evaluate.load("f1") # ============================= # Data # ============================= data_files = {"train": TRAIN_PATH, "validation": VAL_PATH} raw = load_dataset("csv", data_files=data_files) # ensure labels are ints def cast_label(example): example["label"] = int(example["label"]) return example raw = raw.map(cast_label) # Tokenizer try: tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=True) except Exception: tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) def tokenize(batch): return tokenizer( batch["text"], truncation=True, padding=False, max_length=MAX_LEN, ) # tokenize and keep only model inputs + labels tok = raw.map(tokenize, batched=True) # Rename for HF Trainer if "label" in tok["train"].column_names: tok = tok.rename_column("label", "labels") # Remove non-input columns safely (token_type_ids may not exist for some models) keep = {"input_ids", "attention_mask", "token_type_ids", "labels"} cols = tok["train"].column_names remove_cols = [c for c in cols if c not in keep] if remove_cols: tok = tok.remove_columns(remove_cols) # Set PyTorch format tok.set_format(type="torch") train_ds = tok["train"] val_ds = tok["validation"] # Collator data_collator = DataCollatorWithPadding(tokenizer=tokenizer) # ============================= # Class Weights (robust) # ============================= # Derive labels from dataset (ensure 0/1 only) y_train = train_ds["labels"].detach().cpu().numpy() if hasattr(train_ds["labels"], "detach") else np.array(train_ds["labels"]) # type: ignore unique = np.unique(y_train) # Guardrail: map any {0,1,2} style into binary if needed if set(unique) - {0, 1}: # Conservative mapping: {1,2}->1, 0->0 def to_bin(v: int) -> int: return 0 if int(v) == 0 else 1 raw = raw.map(lambda ex: {"label": to_bin(int(ex["label"]))}) tok = raw.map(tokenize, batched=True) if "label" in tok["train"].column_names: tok = tok.rename_column("label", "labels") cols = tok["train"].column_names remove_cols = [c for c in cols if c not in keep] if remove_cols: tok = tok.remove_columns(remove_cols) tok.set_format(type="torch") train_ds = tok["train"] val_ds = tok["validation"] y_train = train_ds["labels"].detach().cpu().numpy() if hasattr(train_ds["labels"], "detach") else np.array(train_ds["labels"]) # type: ignore unique = np.unique(y_train) assert set(unique) <= {0, 1}, f"Found unexpected labels: {unique} — ensure CSVs are binary 0/1." # Environment override: CLASS_WEIGHTS="1.0,3.0" CW_ENV = os.getenv("CLASS_WEIGHTS", "") if CW_ENV: cw = np.array([float(x) for x in CW_ENV.split(",")], dtype=np.float32) assert cw.shape[0] == 2, "CLASS_WEIGHTS must have 2 values for binary classification." else: cw = compute_class_weight(class_weight="balanced", classes=np.array([0, 1]), y=y_train).astype(np.float32) class_weights_tensor = torch.tensor(cw, dtype=torch.float32) # ============================= # Model # ============================= torch.set_num_threads(max(1, (os.cpu_count() or 2) // 2)) model = AutoModelForSequenceClassification.from_pretrained( MODEL_NAME, num_labels=2, id2label=id2label, label2id={k: v for v, k in id2label.items()}, ) # ============================= # Weighted Trainer # ============================= class WeightedTrainer(Trainer): def __init__(self, *args, class_weights: Optional[torch.Tensor] = None, **kwargs): super().__init__(*args, **kwargs) self.class_weights = class_weights self._ce_loss: Optional[nn.Module] = None def compute_loss(self, model, inputs, return_outputs=False, **kwargs): labels = inputs.pop("labels", None) outputs = model(**inputs) logits = outputs.get("logits") if labels is None: loss = outputs["loss"] if "loss" in outputs else None else: if self._ce_loss is None: if self.class_weights is not None: self._ce_loss = nn.CrossEntropyLoss(weight=self.class_weights.to(model.device)) else: self._ce_loss = nn.CrossEntropyLoss() if labels.dtype != torch.long: labels = labels.to(torch.long) loss = self._ce_loss(logits, labels) return (loss, outputs) if return_outputs else loss def compute_metrics(eval_pred): logits, labels = eval_pred preds = np.argmax(logits, axis=1) return { "accuracy": accuracy.compute(predictions=preds, references=labels)["accuracy"], "precision": precision.compute(predictions=preds, references=labels, average="binary")["precision"], "recall": recall.compute(predictions=preds, references=labels, average="binary")["recall"], "f1": f1.compute(predictions=preds, references=labels, average="binary")["f1"], } # ============================= # Training Args # ============================= args = TrainingArguments( output_dir="runs", eval_strategy="epoch", # <-- correct key save_strategy="epoch", logging_strategy="steps", logging_steps=50, per_device_train_batch_size=16, per_device_eval_batch_size=32, gradient_accumulation_steps=2, num_train_epochs=5, learning_rate=3e-5, warmup_ratio=0.1, weight_decay=0.01, load_best_model_at_end=True, metric_for_best_model="f1", greater_is_better=True, report_to=[], # disable W&B et al. seed=42, dataloader_pin_memory=False, push_to_hub=bool(HUB_REPO), hub_model_id=HUB_REPO, hub_private_repo=False, ) trainer = WeightedTrainer( model=model, args=args, train_dataset=train_ds, eval_dataset=val_ds, tokenizer=tokenizer, data_collator=data_collator, class_weights=class_weights_tensor, compute_metrics=compute_metrics, ) trainer.train() metrics = trainer.evaluate() print("Eval:", metrics) trainer.save_model("model") try: tokenizer.save_pretrained("model") except Exception: pass if HUB_REPO: try: trainer.push_to_hub() tokenizer.push_to_hub(HUB_REPO) print(f"Pushed model to {HUB_REPO}") except Exception as e: print(f"Skipping hub push: {e}") else: print("No hub repo specified, model not pushed.")