File size: 4,432 Bytes
3fb84d7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | """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
# The archived worker computes softmax after converting logits to float64.
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()
|