chore(semantics): add v6 release gate pipeline
- Add reproducible v6 boundary, blessing, and consensus-adjudication corpora, plus the tiny-transformer trainer and v6 release-gate evaluator that gate every candidate on the deployed baselines. - Wire consensus-label merging, product-policy anchor evaluation, and sealed blessing benchmark review with their pytest coverage. - Refresh open-training corpus generation, iterative retraining runner, and random-holdout evaluation so v6 candidates can be benchmarked end-to-end.
This commit is contained in:
Rocky
@@ -0,0 +1,491 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Fine-tune bilingual Tiny Transformer challengers for taxonomy-v6 intents."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import random
|
||||
import resource
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.utils.data import DataLoader, Dataset
|
||||
from transformers import AutoModelForSequenceClassification, AutoTokenizer
|
||||
|
||||
|
||||
INTENTS = (
|
||||
"task",
|
||||
"question",
|
||||
"invitation",
|
||||
"complaint",
|
||||
"scheduleNegotiation",
|
||||
"confirmationDecision",
|
||||
"followUpReminder",
|
||||
"blessing",
|
||||
"replyableMessage",
|
||||
"assistantCommand",
|
||||
"informationQuery",
|
||||
"systemNotification",
|
||||
)
|
||||
LANGUAGE_MODELS = {
|
||||
"en": "en",
|
||||
"zh-Hans": "zh",
|
||||
}
|
||||
|
||||
|
||||
def parse_arguments() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--corpus", type=Path, required=True)
|
||||
parser.add_argument("--output-directory", type=Path, required=True)
|
||||
parser.add_argument("--english-model", type=Path, required=True)
|
||||
parser.add_argument("--chinese-model", type=Path, required=True)
|
||||
parser.add_argument("--epochs", type=int, default=3)
|
||||
parser.add_argument("--batch-size", type=int, default=128)
|
||||
parser.add_argument("--max-length", type=int, default=96)
|
||||
parser.add_argument("--learning-rate", type=float, default=3e-4)
|
||||
parser.add_argument("--seed", type=int, default=20260828)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def set_seed(seed: int) -> None:
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
if torch.backends.mps.is_available():
|
||||
torch.mps.manual_seed(seed)
|
||||
|
||||
|
||||
def intent_value(record: dict, intent: str) -> bool:
|
||||
if intent == "replyableMessage":
|
||||
return bool(record.get(intent, record.get("replyable", False)))
|
||||
return bool(record.get(intent, False))
|
||||
|
||||
|
||||
def intent_mask(record: dict) -> list[float]:
|
||||
known_labels = record.get("knownLabels")
|
||||
if known_labels is None:
|
||||
return [1.0] * len(INTENTS)
|
||||
known = set(known_labels)
|
||||
return [float(intent in known) for intent in INTENTS]
|
||||
|
||||
|
||||
def load_records(path: Path) -> dict[str, dict[str, list[dict]]]:
|
||||
records = {
|
||||
language: {"train": [], "validation": [], "test": [], "golden": []}
|
||||
for language in LANGUAGE_MODELS
|
||||
}
|
||||
with path.open(encoding="utf-8") as stream:
|
||||
for line in stream:
|
||||
record = json.loads(line)
|
||||
language = record.get("language")
|
||||
split = record.get("split")
|
||||
if language in records and split in records[language]:
|
||||
records[language][split].append(record)
|
||||
return records
|
||||
|
||||
|
||||
class IntentDataset(Dataset):
|
||||
def __init__(
|
||||
self,
|
||||
records: list[dict],
|
||||
tokenizer: AutoTokenizer,
|
||||
max_length: int,
|
||||
) -> None:
|
||||
encoded = tokenizer(
|
||||
[record["text"] for record in records],
|
||||
max_length=max_length,
|
||||
padding="max_length",
|
||||
truncation=True,
|
||||
return_tensors="pt",
|
||||
)
|
||||
self.inputs = dict(encoded)
|
||||
self.labels = torch.tensor(
|
||||
[
|
||||
[float(intent_value(record, intent)) for intent in INTENTS]
|
||||
for record in records
|
||||
],
|
||||
dtype=torch.float32,
|
||||
)
|
||||
self.masks = torch.tensor(
|
||||
[intent_mask(record) for record in records],
|
||||
dtype=torch.float32,
|
||||
)
|
||||
self.weights = torch.tensor(
|
||||
[float(record.get("sampleWeight", 1.0)) for record in records],
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
def __len__(self) -> int:
|
||||
return self.labels.shape[0]
|
||||
|
||||
def __getitem__(self, index: int) -> dict[str, torch.Tensor]:
|
||||
item = {key: value[index] for key, value in self.inputs.items()}
|
||||
item["labels"] = self.labels[index]
|
||||
item["masks"] = self.masks[index]
|
||||
item["weights"] = self.weights[index]
|
||||
return item
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Metrics:
|
||||
true_positive: int
|
||||
true_negative: int
|
||||
false_positive: int
|
||||
false_negative: int
|
||||
precision: float
|
||||
recall: float
|
||||
f1: float
|
||||
|
||||
def as_dict(self) -> dict:
|
||||
return {
|
||||
"truePositive": self.true_positive,
|
||||
"trueNegative": self.true_negative,
|
||||
"falsePositive": self.false_positive,
|
||||
"falseNegative": self.false_negative,
|
||||
"precision": round(self.precision, 6),
|
||||
"recall": round(self.recall, 6),
|
||||
"f1": round(self.f1, 6),
|
||||
}
|
||||
|
||||
|
||||
def calculate_metrics(
|
||||
expected: np.ndarray,
|
||||
predicted: np.ndarray,
|
||||
mask: np.ndarray,
|
||||
) -> Metrics:
|
||||
expected = expected[mask.astype(bool)].astype(bool)
|
||||
predicted = predicted[mask.astype(bool)].astype(bool)
|
||||
true_positive = int(np.sum(expected & predicted))
|
||||
true_negative = int(np.sum(~expected & ~predicted))
|
||||
false_positive = int(np.sum(~expected & predicted))
|
||||
false_negative = int(np.sum(expected & ~predicted))
|
||||
precision_denominator = true_positive + false_positive
|
||||
recall_denominator = true_positive + false_negative
|
||||
precision = (
|
||||
true_positive / precision_denominator if precision_denominator else 0.0
|
||||
)
|
||||
recall = true_positive / recall_denominator if recall_denominator else 0.0
|
||||
f1 = (
|
||||
2 * precision * recall / (precision + recall)
|
||||
if precision + recall
|
||||
else 0.0
|
||||
)
|
||||
return Metrics(
|
||||
true_positive,
|
||||
true_negative,
|
||||
false_positive,
|
||||
false_negative,
|
||||
precision,
|
||||
recall,
|
||||
f1,
|
||||
)
|
||||
|
||||
|
||||
def predict(
|
||||
model: AutoModelForSequenceClassification,
|
||||
tokenizer: AutoTokenizer,
|
||||
records: list[dict],
|
||||
device: torch.device,
|
||||
max_length: int,
|
||||
batch_size: int,
|
||||
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
|
||||
model.eval()
|
||||
probabilities: list[np.ndarray] = []
|
||||
expected: list[np.ndarray] = []
|
||||
masks: list[np.ndarray] = []
|
||||
with torch.inference_mode():
|
||||
for start in range(0, len(records), batch_size):
|
||||
batch = records[start : start + batch_size]
|
||||
encoded = tokenizer(
|
||||
[record["text"] for record in batch],
|
||||
max_length=max_length,
|
||||
padding="max_length",
|
||||
truncation=True,
|
||||
return_tensors="pt",
|
||||
)
|
||||
inputs = {key: value.to(device) for key, value in encoded.items()}
|
||||
probabilities.append(
|
||||
torch.sigmoid(model(**inputs).logits).cpu().numpy()
|
||||
)
|
||||
expected.append(
|
||||
np.array(
|
||||
[
|
||||
[float(intent_value(record, intent)) for intent in INTENTS]
|
||||
for record in batch
|
||||
],
|
||||
dtype=np.float32,
|
||||
)
|
||||
)
|
||||
masks.append(
|
||||
np.array([intent_mask(record) for record in batch], dtype=np.float32)
|
||||
)
|
||||
return (
|
||||
np.concatenate(probabilities),
|
||||
np.concatenate(expected),
|
||||
np.concatenate(masks),
|
||||
)
|
||||
|
||||
|
||||
def choose_thresholds(
|
||||
probabilities: np.ndarray,
|
||||
expected: np.ndarray,
|
||||
masks: np.ndarray,
|
||||
) -> np.ndarray:
|
||||
thresholds: list[float] = []
|
||||
for index in range(len(INTENTS)):
|
||||
known = masks[:, index].astype(bool)
|
||||
labels = expected[known, index]
|
||||
if not np.any(labels == 1) or not np.any(labels == 0):
|
||||
thresholds.append(0.5)
|
||||
continue
|
||||
candidates = []
|
||||
for threshold in np.linspace(0.05, 0.95, 91):
|
||||
metrics = calculate_metrics(
|
||||
expected[:, index],
|
||||
probabilities[:, index] >= threshold,
|
||||
masks[:, index],
|
||||
)
|
||||
candidates.append((metrics.f1, metrics.precision, metrics.recall, threshold))
|
||||
thresholds.append(float(max(candidates)[3]))
|
||||
return np.array(thresholds, dtype=np.float32)
|
||||
|
||||
|
||||
def summarize(
|
||||
probabilities: np.ndarray,
|
||||
expected: np.ndarray,
|
||||
masks: np.ndarray,
|
||||
thresholds: np.ndarray,
|
||||
) -> dict:
|
||||
per_intent = {}
|
||||
supported_metrics = []
|
||||
for index, intent in enumerate(INTENTS):
|
||||
metrics = calculate_metrics(
|
||||
expected[:, index],
|
||||
probabilities[:, index] >= thresholds[index],
|
||||
masks[:, index],
|
||||
)
|
||||
known_count = int(np.sum(masks[:, index]))
|
||||
positive_count = int(np.sum(expected[:, index] * masks[:, index]))
|
||||
per_intent[intent] = {
|
||||
"knownCount": known_count,
|
||||
"positiveCount": positive_count,
|
||||
**metrics.as_dict(),
|
||||
}
|
||||
if positive_count > 0:
|
||||
supported_metrics.append(metrics)
|
||||
return {
|
||||
"records": int(expected.shape[0]),
|
||||
"evaluatedIntentCount": len(supported_metrics),
|
||||
"macroPrecision": round(
|
||||
float(np.mean([metrics.precision for metrics in supported_metrics])), 6
|
||||
),
|
||||
"macroRecall": round(
|
||||
float(np.mean([metrics.recall for metrics in supported_metrics])), 6
|
||||
),
|
||||
"macroF1": round(
|
||||
float(np.mean([metrics.f1 for metrics in supported_metrics])), 6
|
||||
),
|
||||
"perIntent": per_intent,
|
||||
}
|
||||
|
||||
|
||||
def train_language(
|
||||
language: str,
|
||||
model_path: Path,
|
||||
records: dict[str, list[dict]],
|
||||
arguments: argparse.Namespace,
|
||||
device: torch.device,
|
||||
) -> dict:
|
||||
set_seed(arguments.seed)
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_path, local_files_only=True)
|
||||
model = AutoModelForSequenceClassification.from_pretrained(
|
||||
model_path,
|
||||
local_files_only=True,
|
||||
num_labels=len(INTENTS),
|
||||
problem_type="multi_label_classification",
|
||||
ignore_mismatched_sizes=True,
|
||||
).to(device)
|
||||
dataset = IntentDataset(records["train"], tokenizer, arguments.max_length)
|
||||
generator = torch.Generator().manual_seed(arguments.seed)
|
||||
loader = DataLoader(
|
||||
dataset,
|
||||
batch_size=arguments.batch_size,
|
||||
shuffle=True,
|
||||
generator=generator,
|
||||
)
|
||||
weighted_positive = (dataset.labels * dataset.masks) * dataset.weights[:, None]
|
||||
weighted_known = dataset.masks * dataset.weights[:, None]
|
||||
positive_counts = weighted_positive.sum(dim=0)
|
||||
negative_counts = weighted_known.sum(dim=0) - positive_counts
|
||||
positive_weights = torch.clamp(
|
||||
negative_counts / torch.clamp(positive_counts, min=1),
|
||||
min=1,
|
||||
max=20,
|
||||
).to(device)
|
||||
criterion = torch.nn.BCEWithLogitsLoss(
|
||||
pos_weight=positive_weights,
|
||||
reduction="none",
|
||||
)
|
||||
optimizer = torch.optim.AdamW(
|
||||
model.parameters(),
|
||||
lr=arguments.learning_rate,
|
||||
weight_decay=0.01,
|
||||
)
|
||||
|
||||
epoch_losses = []
|
||||
training_started = time.perf_counter()
|
||||
for epoch in range(arguments.epochs):
|
||||
model.train()
|
||||
running_loss = 0.0
|
||||
for batch in loader:
|
||||
labels = batch.pop("labels").to(device)
|
||||
masks = batch.pop("masks").to(device)
|
||||
weights = batch.pop("weights").to(device).unsqueeze(1)
|
||||
inputs = {key: value.to(device) for key, value in batch.items()}
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
losses = criterion(model(**inputs).logits, labels)
|
||||
weighted_masks = masks * weights
|
||||
loss = (losses * weighted_masks).sum() / weighted_masks.sum().clamp(min=1)
|
||||
loss.backward()
|
||||
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
|
||||
optimizer.step()
|
||||
running_loss += float(loss.detach().cpu())
|
||||
average_loss = running_loss / max(len(loader), 1)
|
||||
epoch_losses.append(round(average_loss, 6))
|
||||
print(
|
||||
f"TRAIN language={language} epoch={epoch + 1}/{arguments.epochs} "
|
||||
f"loss={average_loss:.6f}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
predictions = {}
|
||||
for split in ("validation", "test", "golden"):
|
||||
predictions[split] = predict(
|
||||
model,
|
||||
tokenizer,
|
||||
records[split],
|
||||
device,
|
||||
arguments.max_length,
|
||||
arguments.batch_size,
|
||||
)
|
||||
thresholds = choose_thresholds(*predictions["validation"])
|
||||
evaluations = {
|
||||
split: summarize(*values, thresholds)
|
||||
for split, values in predictions.items()
|
||||
}
|
||||
|
||||
sample_text = records["test"][0]["text"]
|
||||
encoded = tokenizer(
|
||||
sample_text,
|
||||
max_length=arguments.max_length,
|
||||
padding="max_length",
|
||||
truncation=True,
|
||||
return_tensors="pt",
|
||||
)
|
||||
inputs = {key: value.to(device) for key, value in encoded.items()}
|
||||
model.eval()
|
||||
with torch.inference_mode():
|
||||
started = time.perf_counter()
|
||||
model(**inputs)
|
||||
if device.type == "mps":
|
||||
torch.mps.synchronize()
|
||||
cold_ms = (time.perf_counter() - started) * 1_000
|
||||
warm_samples = []
|
||||
for _ in range(100):
|
||||
started = time.perf_counter()
|
||||
model(**inputs)
|
||||
if device.type == "mps":
|
||||
torch.mps.synchronize()
|
||||
warm_samples.append((time.perf_counter() - started) * 1_000)
|
||||
|
||||
output = arguments.output_directory / LANGUAGE_MODELS[language]
|
||||
output.mkdir(parents=True, exist_ok=True)
|
||||
model.save_pretrained(output)
|
||||
tokenizer.save_pretrained(output)
|
||||
model_bytes = sum(path.stat().st_size for path in output.iterdir() if path.is_file())
|
||||
return {
|
||||
"language": language,
|
||||
"baseModel": str(model_path),
|
||||
"trainRecords": len(records["train"]),
|
||||
"epochLosses": epoch_losses,
|
||||
"trainingSeconds": round(time.perf_counter() - training_started, 3),
|
||||
"thresholds": {
|
||||
intent: round(float(thresholds[index]), 4)
|
||||
for index, intent in enumerate(INTENTS)
|
||||
},
|
||||
"evaluations": evaluations,
|
||||
"runtime": {
|
||||
"engine": f"PyTorch eager on {device.type}",
|
||||
"coldMilliseconds": round(cold_ms, 3),
|
||||
"warmMeanMilliseconds": round(float(np.mean(warm_samples)), 3),
|
||||
"warmP95Milliseconds": round(float(np.percentile(warm_samples, 95)), 3),
|
||||
"processMaximumRSSBytes": int(
|
||||
resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
|
||||
),
|
||||
},
|
||||
"savedModelBytes": model_bytes,
|
||||
}
|
||||
|
||||
|
||||
def main() -> None:
|
||||
arguments = parse_arguments()
|
||||
set_seed(arguments.seed)
|
||||
device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
|
||||
records = load_records(arguments.corpus)
|
||||
model_paths = {
|
||||
"en": arguments.english_model,
|
||||
"zh-Hans": arguments.chinese_model,
|
||||
}
|
||||
language_reports = []
|
||||
for language in ("zh-Hans", "en"):
|
||||
counts = {
|
||||
split: len(split_records)
|
||||
for split, split_records in records[language].items()
|
||||
}
|
||||
print(f"DATA language={language} counts={counts}", flush=True)
|
||||
language_reports.append(
|
||||
train_language(
|
||||
language,
|
||||
model_paths[language],
|
||||
records[language],
|
||||
arguments,
|
||||
device,
|
||||
)
|
||||
)
|
||||
if device.type == "mps":
|
||||
torch.mps.empty_cache()
|
||||
|
||||
report = {
|
||||
"schemaVersion": 1,
|
||||
"purpose": "Taxonomy-v6 Tiny Transformer research challenger",
|
||||
"corpus": str(arguments.corpus),
|
||||
"seed": arguments.seed,
|
||||
"intents": list(INTENTS),
|
||||
"configuration": {
|
||||
"epochs": arguments.epochs,
|
||||
"batchSize": arguments.batch_size,
|
||||
"maxLength": arguments.max_length,
|
||||
"learningRate": arguments.learning_rate,
|
||||
},
|
||||
"languages": language_reports,
|
||||
"limitations": [
|
||||
"Thresholds use only the frozen validation split, which has 20 records per language.",
|
||||
"PyTorch runtime is not directly comparable with Core ML runtime.",
|
||||
"The Chinese UER checkpoint does not declare a model-weight license in its model card.",
|
||||
],
|
||||
}
|
||||
arguments.output_directory.mkdir(parents=True, exist_ok=True)
|
||||
report_path = arguments.output_directory / "training-evaluation-report.json"
|
||||
report_path.write_text(
|
||||
json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True) + "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
print(f"REPORT {report_path}", flush=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user