| """Load SafeAge Guard as a binary classifier over its pinned Qwen base.""" |
| import argparse |
| import json |
| from pathlib import Path |
|
|
| DEFAULT_REPO = "mehedi4/SafeAge-Guard" |
|
|
|
|
| def load_guard(source=DEFAULT_REPO, revision=None, device=None): |
| import torch |
| from huggingface_hub import snapshot_download |
| from peft import PeftModel |
| from transformers import AutoModelForSequenceClassification, AutoTokenizer |
|
|
| folder = Path(source) |
| if not folder.is_dir(): |
| folder = Path(snapshot_download( |
| repo_id=source, revision=revision, |
| allow_patterns=["adapter_config.json", "adapter_model.safetensors", |
| "base_model_reference.json", "config.json", "decision_threshold.json", |
| "tokenizer.json", "tokenizer_config.json", "chat_template.jinja"], |
| )) |
| 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.padding_side = "right" |
| tokenizer.truncation_side = "right" |
| tokenizer.model_max_length = reference["max_length"] |
| if tokenizer.pad_token_id != reference["pad_token_id"]: |
| raise ValueError("Tokenizer padding token differs from the original checkpoint") |
| labels = {int(k): v for k, v in reference["id2label"].items()} |
| base = AutoModelForSequenceClassification.from_pretrained( |
| reference["model_source"], revision=reference["model_revision"], |
| num_labels=reference["num_labels"], id2label=labels, |
| label2id={v: k for k, v in labels.items()}, |
| use_safetensors=True, attn_implementation="sdpa", dtype=torch.float32, |
| ) |
| base.config.pad_token_id = reference["pad_token_id"] |
| base.config.problem_type = "single_label_classification" |
| base.config.use_cache = False |
| model = PeftModel.from_pretrained( |
| base, folder, is_trainable=False, local_files_only=True, |
| ).eval() |
| model.requires_grad_(False) |
| device = device or ("cuda" if torch.cuda.is_available() else "cpu") |
| model.to(device) |
| return model, tokenizer, reference, decision |
|
|
|
|
| def classify_texts(texts, guard): |
| """Classify raw requests; do not apply a chat template.""" |
| 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 = guard |
| encoded = tokenizer( |
| texts, truncation=True, max_length=reference["max_length"], |
| padding=True, pad_to_multiple_of=8, return_tensors="pt", |
| ) |
| device = next(model.parameters()).device |
| encoded = {key: value.to(device) for key, value in encoded.items()} |
| with torch.inference_mode(): |
| probabilities = torch.softmax(model(**encoded).logits.float(), dim=-1) |
| values = probabilities[:, 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 with the 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 model folder") |
| parser.add_argument("--revision", help="Optional Hub commit to pin the adapter release") |
| parser.add_argument("--device", help="For example cpu or cuda; default uses CUDA when available") |
| args = parser.parse_args() |
| text = args.text if args.text is not None else args.text_file.read_text(encoding="utf-8") |
| guard = load_guard(source=args.source, revision=args.revision, device=args.device) |
| print(json.dumps(classify_texts(text, guard)[0], ensure_ascii=False, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
