| """Load the original fully fine-tuned SafeAge XLM-R binary baseline.""" |
| import argparse |
| import contextlib |
| import json |
| from pathlib import Path |
|
|
| DEFAULT_REPO = "mehedi4/SafeAge-XLM-R-Base" |
|
|
|
|
| def load_classifier(source=DEFAULT_REPO, revision=None, device=None, precision="auto"): |
| import torch |
| from huggingface_hub import snapshot_download |
| from transformers import AutoModelForSequenceClassification, AutoTokenizer |
|
|
| folder = Path(source) |
| if not folder.is_dir(): |
| folder = Path(snapshot_download( |
| repo_id=source, revision=revision, |
| allow_patterns=["config.json", "model.safetensors", "decision_threshold.json", |
| "tokenizer.json", "tokenizer_config.json", "base_model_reference.json"], |
| )) |
| reference = json.loads((folder / "base_model_reference.json").read_text()) |
| decision = json.loads((folder / "decision_threshold.json").read_text()) |
| tokenizer = AutoTokenizer.from_pretrained(folder, use_fast=True, local_files_only=True) |
| tokenizer.model_max_length = reference["max_length"] |
| tokenizer.padding_side = "right" |
| tokenizer.truncation_side = "right" |
| if tokenizer.pad_token_id != reference["pad_token_id"]: |
| raise ValueError("Padding token differs from the original evaluated tokenizer") |
| model = AutoModelForSequenceClassification.from_pretrained( |
| folder, local_files_only=True, use_safetensors=True, dtype=torch.float32, |
| ).eval() |
| if model.config.model_type != "xlm-roberta" or model.config.num_labels != 2: |
| raise ValueError("Expected the original binary XLM-R classifier") |
| model.requires_grad_(False) |
| device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu")) |
| if precision not in {"auto", "fp32", "fp16_mixed"}: |
| raise ValueError("Precision must be auto, fp32, or fp16_mixed") |
| mixed = device.type == "cuda" if precision == "auto" else precision == "fp16_mixed" |
| if mixed and device.type != "cuda": |
| raise ValueError("FP16 mixed inference requires a CUDA device; use fp32 on CPU") |
| model.to(device) |
| return model, tokenizer, reference, decision, mixed |
|
|
|
|
| def classify_texts(texts, classifier): |
| """Classify complete raw requests with context; no chat template is applied.""" |
| import torch |
| if isinstance(texts, str): |
| texts = [texts] |
| if not texts or any(not isinstance(text, str) for text in texts): |
| raise ValueError("Supply one or more strings") |
| model, tokenizer, reference, decision, mixed = classifier |
| encoded = tokenizer( |
| texts, truncation=True, max_length=reference["max_length"], padding=True, |
| pad_to_multiple_of=8, return_token_type_ids=False, return_tensors="pt", |
| ) |
| device = next(model.parameters()).device |
| encoded = {key: value.to(device) for key, value in encoded.items()} |
| amp = torch.autocast(device_type="cuda", dtype=torch.float16) if mixed else contextlib.nullcontext() |
| with torch.inference_mode(), amp: |
| logits = model(**encoded).logits |
| |
| values = torch.softmax(logits.double(), dim=-1)[:, decision["positive_class_id"]].cpu().tolist() |
| threshold = float(decision["threshold"]) |
| return [{"decision": "block" if p >= threshold else "allow", |
| "block_probability": p, "threshold": threshold} for p in values] |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| group = parser.add_mutually_exclusive_group(required=True) |
| group.add_argument("--text", help="Complete raw request and its visible context") |
| group.add_argument("--text-file", type=Path, help="UTF-8 file containing one complete request") |
| parser.add_argument("--source", default=DEFAULT_REPO, help="Hub repository or local folder") |
| parser.add_argument("--revision", help="Optional exact release commit") |
| parser.add_argument("--device", help="cpu, cuda, or cuda:0; default selects CUDA if available") |
| parser.add_argument("--precision", choices=["auto", "fp32", "fp16_mixed"], default="auto") |
| args = parser.parse_args() |
| text = args.text if args.text is not None else args.text_file.read_text(encoding="utf-8") |
| classifier = load_classifier(args.source, args.revision, args.device, args.precision) |
| print(json.dumps(classify_texts(text, classifier)[0], ensure_ascii=False, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
