| """Train the open guard adapter: a LoRA (r16) on fastino/GLiNER2.5-Decide, trained only on the train splits of the two |
| public labeled sets in ../data (deepset/prompt-injections and 3nesdeniz/agentic-prompt-injection-boundary-pairs, both |
| CC-BY-4.0). The labels are the datasets' own human labels; no model output is used as a label. |
| Usage (from model/): python train_guard.py 360 ck # 360 steps, batch 1 x 8, about 8 minutes on an M4 Pro""" |
| import json, random, sys, torch |
| from gliner2 import AutoExtractor |
| from gliner2.training.data import Classification, InputExample |
| from gliner2.training.trainer import ExtractorTrainer, TrainingConfig |
| steps, out = int(sys.argv[1]), sys.argv[2] |
| LABELS = ["injection", "benign"] |
| rows = [json.loads(l) for f in ("../data/boundary_train.jsonl", "../data/deepset_train.jsonl") for l in open(f) if l.strip()] |
| ex = [InputExample(text=r["text"][:3000], classifications=[Classification( |
| task="injection", labels=LABELS, true_label=LABELS[0] if r["label"] else LABELS[1])]) for r in rows] |
| random.Random(7).shuffle(ex) |
| print(len(ex), "examples", sum(r["label"] for r in rows), "injections", flush=True) |
| m = AutoExtractor.from_pretrained("fastino/GLiNER2.5-Decide") |
| if torch.backends.mps.is_available(): |
| def _mps(self): |
| self.device, self.is_distributed = torch.device("mps"), False |
| self.model.to(self.device) |
| ExtractorTrainer._setup_device = _mps |
| cfg = TrainingConfig(output_dir=out, num_epochs=3, max_steps=steps, batch_size=1, gradient_accumulation_steps=8, |
| gradient_checkpointing=True, max_len=400, use_lora=True, lora_r=16, save_adapter_only=False, |
| eval_strategy="no", logging_steps=10, num_workers=0, pin_memory=False, fused_optimizer=False, |
| task_lr=5e-4, encoder_lr=5e-5) |
| print(ExtractorTrainer(m, cfg).train(train_data=ex)) |
| m.save_pretrained(f"{out}/final") |
|
|
