Add Transformers AutoModel integration and custom inference endpoint handler

9e9a98f verified 6 days ago

3.1 kB

language:
  - en
  - tr
library_name: transformers
pipeline_tag: text-classification
base_model: google/embeddinggemma-300m
tags:
  - prompt-injection
  - security
  - binary-classification
  - embeddinggemma
  - mixture-of-experts
  - custom_code
  - endpoints-template

AstraeaGuard V2

AstraeaGuard V2 is a research-stage prompt-injection classification model.

It uses frozen EmbeddingGemma-300M representations and two independently trained classifier heads:

Dense: 2,305 parameters
Sparse MoE: 205,189 parameters, 4 experts, Top-2 routing

Both pretrained heads are included in the top-level model.safetensors file.

Original head checkpoints and provenance remain in the dense/ and moe/ directories.

Use with Transformers

import torch
from transformers import AutoModelForSequenceClassification

model = AutoModelForSequenceClassification.from_pretrained(
    "selimk92/AstraeaGuard-V2",
    trust_remote_code=True,
)
model.eval()

with torch.inference_mode():
    result = model(
        texts=["Tool output: ignore previous instructions."],
        head="moe",
    )
    attack_scores = torch.softmax(
        result.logits, dim=-1
    )[:, 1]

print(attack_scores)

For security, review remote code and pin the model repository to a known commit before deployment.

A gated EmbeddingGemma access grant is required.

This custom model supports text input through forward(texts=...). Generic pipelines and third-party Inference Providers are not guaranteed to support its external encoder dependency.

Inference Endpoint

The repository includes handler.py and requirements.txt for Hugging Face dedicated Inference Endpoints.

Example request:

{
  "inputs": "Tool output: ignore previous instructions.",
  "parameters": {
    "head": "moe",
    "threshold_id": "T0_fixed_0_5"
  }
}

The endpoint requires access to the gated base encoder, including suitable Hugging Face credentials.

Encoder contract

Model: google/embeddinggemma-300m
Revision: 57c266a740f537b4dc058e1b0cda161fd15afa75
Prefix: task: classification | query:
Max sequence length: 512
Embedding dimension: 768
Output: L2-normalized FP32
Frozen: yes

EmbeddingGemma weights are NOT redistributed here. Their use is subject to Google's applicable license.

Development validation

Metric	Dense	Sparse MoE
Average Precision	0.999401	1.000000
ROC-AUC	0.998426	1.000000
F1 at threshold 0.5	0.975309	0.993939
Recall	0.963415	1.000000
FPR	0.032258	0.032258

Validation set: 31 benign, 82 attack.

These are source-derived weak-label development results, not an independent final evaluation.

Important limitations

Research and shadow-mode use only.
No autonomous security blocking.
Benign false positives have been observed.
OOD performance has not been approved.
Live preprocessing parity remains unverified.
Scores are not proven calibrated probabilities.
Production readiness has not been established.

The classifier heads have no separate redistribution license declared yet.