Nomearod Claude Opus 4.6 (1M context) commited on
Commit
2dc5235
·
1 Parent(s): cebf463

feat(security): add Modal DeBERTa injection classifier deployment

Browse files

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Files changed (1) hide show
  1. modal/injection_classifier.py +59 -0
modal/injection_classifier.py ADDED
@@ -0,0 +1,59 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Deploy DeBERTa-v3-base injection classifier on Modal.
2
+
3
+ Usage:
4
+ modal deploy modal/injection_classifier.py
5
+ modal serve modal/injection_classifier.py # Dev mode
6
+
7
+ Endpoint: POST /classify {"text": "..."}
8
+ Returns: {"label": "INJECTION" | "SAFE", "score": 0.95}
9
+ """
10
+
11
+ import modal
12
+
13
+ MODELS_DIR = "/models"
14
+
15
+ classifier_image = (
16
+ modal.Image.debian_slim(python_version="3.11")
17
+ .pip_install(
18
+ "transformers>=4.40.0",
19
+ "torch>=2.0.0",
20
+ "sentencepiece",
21
+ "protobuf",
22
+ )
23
+ )
24
+
25
+ app = modal.App("agent-bench-injection-classifier")
26
+ model_volume = modal.Volume.from_name("injection-model-cache", create_if_missing=True)
27
+
28
+
29
+ @app.cls(
30
+ image=classifier_image,
31
+ gpu="T4",
32
+ scaledown_window=300,
33
+ timeout=120,
34
+ volumes={MODELS_DIR: model_volume},
35
+ )
36
+ class InjectionClassifier:
37
+ @modal.enter()
38
+ def load(self):
39
+ from transformers import pipeline
40
+
41
+ self.pipe = pipeline(
42
+ "text-classification",
43
+ model="deepset/deberta-v3-base-injection",
44
+ device="cuda",
45
+ model_kwargs={"cache_dir": MODELS_DIR},
46
+ )
47
+
48
+ @modal.method()
49
+ def classify(self, text: str) -> dict:
50
+ result = self.pipe(text, truncation=True, max_length=512)[0]
51
+ return {"label": result["label"], "score": result["score"]}
52
+
53
+
54
+ @app.function(image=classifier_image, gpu="T4", volumes={MODELS_DIR: model_volume})
55
+ @modal.web_endpoint(method="POST")
56
+ def classify_endpoint(item: dict) -> dict:
57
+ """HTTP endpoint wrapper for the classifier."""
58
+ classifier = InjectionClassifier()
59
+ return classifier.classify.remote(item["text"])