| """Türkçe prompt injection ikili sınıflandırıcısını eğitir (encoder, softmax iki sınıf). |
| |
| Kullanım: |
| python -m pinj.train --base models/berturk-cased --data data/pinj --out runs/pinj-berturk |
| Çıktı: runs/.../final (model + tokenizer), runs/.../train_report.json (dev metrikleri) |
| """ |
| import argparse |
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| from torch.utils.data import Dataset |
| from transformers import (AutoModelForSequenceClassification, AutoTokenizer, Trainer, TrainingArguments, |
| set_seed) |
|
|
| ID2LABEL = {0: "BENIGN", 1: "INJECTION"} |
| LABEL2ID = {v: k for k, v in ID2LABEL.items()} |
|
|
|
|
| class PinjData(Dataset): |
| def __init__(self, path, tok, max_len): |
| rows = [json.loads(l) for l in open(path, encoding="utf-8")] |
| self.enc = tok([r["text"] for r in rows], truncation=True, max_length=max_len) |
| self.y = np.array([r["label"] for r in rows], dtype=np.int64) |
|
|
| def __len__(self): |
| return len(self.y) |
|
|
| def __getitem__(self, i): |
| item = {k: v[i] for k, v in self.enc.items()} |
| item["labels"] = self.y[i] |
| return item |
|
|
|
|
| def collate(batch, pad_id): |
| n = max(len(b["input_ids"]) for b in batch) |
| out = {} |
| for k in batch[0]: |
| if k == "labels": |
| out[k] = torch.tensor([b[k] for b in batch]) |
| else: |
| pad = pad_id if k == "input_ids" else 0 |
| out[k] = torch.tensor([list(b[k]) + [pad] * (n - len(b[k])) for b in batch]) |
| return out |
|
|
|
|
| def prf(y, p): |
| tp = int(((p == 1) & (y == 1)).sum()) |
| fp = int(((p == 1) & (y == 0)).sum()) |
| fn = int(((p == 0) & (y == 1)).sum()) |
| pr = tp / (tp + fp) if tp + fp else 0.0 |
| rc = tp / (tp + fn) if tp + fn else 0.0 |
| return pr, rc, (2 * pr * rc / (pr + rc) if pr + rc else 0.0) |
|
|
|
|
| def compute_metrics(pred): |
| p = pred.predictions.argmax(-1) |
| pr, rc, f1 = prf(pred.label_ids, p) |
| return {"f1": f1, "precision": pr, "recall": rc} |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--base", required=True) |
| ap.add_argument("--data", required=True) |
| ap.add_argument("--out", required=True) |
| ap.add_argument("--epochs", type=float, default=4) |
| ap.add_argument("--lr", type=float, default=3e-5) |
| ap.add_argument("--batch", type=int, default=32) |
| ap.add_argument("--max-len", type=int, default=384) |
| ap.add_argument("--seed", type=int, default=42) |
| |
| ap.add_argument("--fp32", action="store_true") |
| args = ap.parse_args() |
|
|
| set_seed(args.seed) |
| tok = AutoTokenizer.from_pretrained(args.base) |
| data = Path(args.data) |
| train = PinjData(data / "train.jsonl", tok, args.max_len) |
| dev = PinjData(data / "dev.jsonl", tok, args.max_len) |
| model = AutoModelForSequenceClassification.from_pretrained( |
| args.base, num_labels=2, id2label=ID2LABEL, label2id=LABEL2ID) |
|
|
| targs = TrainingArguments( |
| output_dir=args.out, num_train_epochs=args.epochs, learning_rate=args.lr, |
| per_device_train_batch_size=args.batch, per_device_eval_batch_size=64, |
| warmup_steps=0.06, weight_decay=0.01, lr_scheduler_type="linear", bf16=not args.fp32, |
| eval_strategy="epoch", save_strategy="epoch", save_total_limit=2, load_best_model_at_end=True, |
| metric_for_best_model="f1", greater_is_better=True, logging_steps=50, report_to="none", |
| seed=args.seed, train_sampling_strategy="group_by_length", |
| ) |
| trainer = Trainer(model=model, args=targs, train_dataset=train, eval_dataset=dev, |
| data_collator=lambda b: collate(b, tok.pad_token_id), compute_metrics=compute_metrics) |
| trainer.train() |
| final = Path(args.out) / "final" |
| trainer.save_model(final) |
| |
| tok.save_pretrained(final) |
| hist = [h for h in trainer.state.log_history if "eval_f1" in h] |
| rep = {"base": args.base, "data": args.data, "args": vars(args), "dev_by_epoch": hist, |
| "best": trainer.state.best_metric} |
| (Path(args.out) / "train_report.json").write_text(json.dumps(rep, indent=1), encoding="utf-8") |
| for h in hist: |
| print(f"epoch {h['epoch']:.0f} dev F1 {h['eval_f1']:.4f} P {h['eval_precision']:.4f} " |
| f"R {h['eval_recall']:.4f} loss {h['eval_loss']:.4f}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
