spamkit-api 0.1.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- spamkit/__init__.py +27 -0
- spamkit/app.py +314 -0
- spamkit/classifier.py +122 -0
- spamkit/cli.py +199 -0
- spamkit/client.py +174 -0
- spamkit/errors.py +40 -0
- spamkit/evaluation.py +180 -0
- spamkit/observability.py +86 -0
- spamkit/schemas.py +121 -0
- spamkit/settings.py +67 -0
- spamkit/verdict.py +46 -0
- spamkit_api-0.1.0.dist-info/METADATA +245 -0
- spamkit_api-0.1.0.dist-info/RECORD +16 -0
- spamkit_api-0.1.0.dist-info/WHEEL +4 -0
- spamkit_api-0.1.0.dist-info/entry_points.txt +2 -0
- spamkit_api-0.1.0.dist-info/licenses/LICENSE +661 -0
spamkit/__init__.py
ADDED
|
@@ -0,0 +1,27 @@
|
|
|
1
|
+
"""Spamkit: spam scoring for user-submitted text.
|
|
2
|
+
|
|
3
|
+
The server (``spamkit.app``) needs the ``server`` extra. The client (``spamkit.client``)
|
|
4
|
+
and the threshold and error types need only the base install.
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
from spamkit.client import (
|
|
8
|
+
SpamkitClient,
|
|
9
|
+
SpamkitError,
|
|
10
|
+
SpamkitTimeoutError,
|
|
11
|
+
SpamkitTransportError,
|
|
12
|
+
)
|
|
13
|
+
from spamkit.errors import ErrorCode
|
|
14
|
+
from spamkit.verdict import Thresholds, Verdict
|
|
15
|
+
|
|
16
|
+
__version__ = "0.1.0"
|
|
17
|
+
|
|
18
|
+
__all__ = [
|
|
19
|
+
"ErrorCode",
|
|
20
|
+
"SpamkitClient",
|
|
21
|
+
"SpamkitError",
|
|
22
|
+
"SpamkitTimeoutError",
|
|
23
|
+
"SpamkitTransportError",
|
|
24
|
+
"Thresholds",
|
|
25
|
+
"Verdict",
|
|
26
|
+
"__version__",
|
|
27
|
+
]
|
spamkit/app.py
ADDED
|
@@ -0,0 +1,314 @@
|
|
|
1
|
+
"""FastAPI application for the v1 scoring API."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import hmac
|
|
6
|
+
import json
|
|
7
|
+
import re
|
|
8
|
+
import threading
|
|
9
|
+
import time
|
|
10
|
+
from collections.abc import AsyncIterator, Sequence
|
|
11
|
+
from contextlib import asynccontextmanager
|
|
12
|
+
from typing import Any
|
|
13
|
+
from uuid import uuid4
|
|
14
|
+
|
|
15
|
+
from fastapi import FastAPI, Request
|
|
16
|
+
from fastapi.responses import JSONResponse, PlainTextResponse
|
|
17
|
+
from starlette.concurrency import run_in_threadpool
|
|
18
|
+
|
|
19
|
+
from spamkit import __version__
|
|
20
|
+
from spamkit.classifier import MODEL_ID, ScoredText, Scorer, SpamtonClassifier
|
|
21
|
+
from spamkit.errors import ApiError, ErrorCode
|
|
22
|
+
from spamkit.observability import Metrics, log_event, logger, score_band
|
|
23
|
+
from spamkit.schemas import (
|
|
24
|
+
API_VERSION,
|
|
25
|
+
BatchItemError,
|
|
26
|
+
BatchItemOk,
|
|
27
|
+
BatchResponse,
|
|
28
|
+
ErrorBody,
|
|
29
|
+
HealthResponse,
|
|
30
|
+
InfoResponse,
|
|
31
|
+
ReadyResponse,
|
|
32
|
+
ScoreInput,
|
|
33
|
+
ScoreResult,
|
|
34
|
+
parse_score_input,
|
|
35
|
+
)
|
|
36
|
+
from spamkit.settings import Settings
|
|
37
|
+
|
|
38
|
+
_REQUEST_ID = re.compile(r"[A-Za-z0-9._-]{1,128}")
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
class _ServiceState:
|
|
42
|
+
def __init__(self, settings: Settings) -> None:
|
|
43
|
+
self.settings = settings
|
|
44
|
+
self.scorer: Scorer | None = None
|
|
45
|
+
self.load_error: str | None = None
|
|
46
|
+
self.metrics = Metrics()
|
|
47
|
+
self._inflight = 0
|
|
48
|
+
self._gate = threading.Lock()
|
|
49
|
+
|
|
50
|
+
def try_acquire(self) -> bool:
|
|
51
|
+
"""Admission control: refuse work when saturated instead of queueing without bound."""
|
|
52
|
+
with self._gate:
|
|
53
|
+
if self._inflight >= self.settings.max_inflight:
|
|
54
|
+
return False
|
|
55
|
+
self._inflight += 1
|
|
56
|
+
return True
|
|
57
|
+
|
|
58
|
+
def release(self) -> None:
|
|
59
|
+
with self._gate:
|
|
60
|
+
self._inflight -= 1
|
|
61
|
+
|
|
62
|
+
|
|
63
|
+
def _state(request: Request) -> _ServiceState:
|
|
64
|
+
return request.app.state.spamkit # type: ignore[no-any-return]
|
|
65
|
+
|
|
66
|
+
|
|
67
|
+
def _require_auth(request: Request, settings: Settings) -> None:
|
|
68
|
+
if settings.api_token is None:
|
|
69
|
+
return
|
|
70
|
+
supplied = request.headers.get("authorization", "").encode()
|
|
71
|
+
expected = f"Bearer {settings.api_token}".encode()
|
|
72
|
+
if not hmac.compare_digest(supplied, expected):
|
|
73
|
+
raise ApiError(
|
|
74
|
+
ErrorCode.UNAUTHORIZED,
|
|
75
|
+
"A valid bearer token is required.",
|
|
76
|
+
headers={"WWW-Authenticate": "Bearer"},
|
|
77
|
+
)
|
|
78
|
+
|
|
79
|
+
|
|
80
|
+
async def _read_json(request: Request, limit: int) -> Any:
|
|
81
|
+
"""Read the body with a running byte count, so an oversized upload is refused as soon as it crosses the limit."""
|
|
82
|
+
declared = request.headers.get("content-length")
|
|
83
|
+
if declared is not None and declared.isdigit() and int(declared) > limit:
|
|
84
|
+
raise ApiError(ErrorCode.INPUT_TOO_LARGE, f"Request body must be at most {limit} bytes.")
|
|
85
|
+
chunks: list[bytes] = []
|
|
86
|
+
total = 0
|
|
87
|
+
async for chunk in request.stream():
|
|
88
|
+
total += len(chunk)
|
|
89
|
+
if total > limit:
|
|
90
|
+
raise ApiError(ErrorCode.INPUT_TOO_LARGE, f"Request body must be at most {limit} bytes.")
|
|
91
|
+
chunks.append(chunk)
|
|
92
|
+
try:
|
|
93
|
+
return json.loads(b"".join(chunks))
|
|
94
|
+
except RecursionError as exc:
|
|
95
|
+
# Very deep nesting exhausts the parser's stack. That is a bad request, not a server fault.
|
|
96
|
+
raise ApiError(ErrorCode.INVALID_REQUEST, "Request body is nested too deeply.") from exc
|
|
97
|
+
except ValueError as exc:
|
|
98
|
+
raise ApiError(ErrorCode.INVALID_REQUEST, "Request body must be valid JSON.") from exc
|
|
99
|
+
|
|
100
|
+
|
|
101
|
+
async def _score_texts(state: _ServiceState, texts: Sequence[str]) -> tuple[Scorer, list[ScoredText], float]:
|
|
102
|
+
"""Score ``texts`` and return the scorer that produced them, so callers never re-read shared state."""
|
|
103
|
+
scorer = state.scorer
|
|
104
|
+
if scorer is None:
|
|
105
|
+
raise ApiError(ErrorCode.MODEL_NOT_READY, "The model is not loaded. Check the startup logs.")
|
|
106
|
+
if not state.try_acquire():
|
|
107
|
+
raise ApiError(
|
|
108
|
+
ErrorCode.RATE_LIMITED,
|
|
109
|
+
"The service is at capacity. Retry shortly.",
|
|
110
|
+
headers={"Retry-After": "1"},
|
|
111
|
+
)
|
|
112
|
+
started = time.perf_counter()
|
|
113
|
+
try:
|
|
114
|
+
scored = await run_in_threadpool(scorer.score_many, list(texts))
|
|
115
|
+
finally:
|
|
116
|
+
state.release()
|
|
117
|
+
return scorer, scored, (time.perf_counter() - started) * 1000
|
|
118
|
+
|
|
119
|
+
|
|
120
|
+
def _build_result(
|
|
121
|
+
state: _ServiceState,
|
|
122
|
+
scorer: Scorer,
|
|
123
|
+
item: ScoreInput,
|
|
124
|
+
scored: ScoredText,
|
|
125
|
+
latency_ms: float,
|
|
126
|
+
request_id: str,
|
|
127
|
+
) -> ScoreResult:
|
|
128
|
+
verdict = state.settings.thresholds.verdict_for(scored.score)
|
|
129
|
+
state.metrics.record_verdict(verdict)
|
|
130
|
+
log_event(
|
|
131
|
+
"score",
|
|
132
|
+
request_id=request_id,
|
|
133
|
+
verdict=verdict.value,
|
|
134
|
+
score_band=score_band(scored.score),
|
|
135
|
+
latency_ms=round(latency_ms, 1),
|
|
136
|
+
model_version=scorer.model_version,
|
|
137
|
+
context=item.context,
|
|
138
|
+
truncated=scored.truncated,
|
|
139
|
+
)
|
|
140
|
+
return ScoreResult(
|
|
141
|
+
id=item.id,
|
|
142
|
+
score=scored.score,
|
|
143
|
+
verdict=verdict,
|
|
144
|
+
model=scorer.model_id,
|
|
145
|
+
model_version=scorer.model_version,
|
|
146
|
+
latency_ms=round(latency_ms),
|
|
147
|
+
truncated=scored.truncated,
|
|
148
|
+
)
|
|
149
|
+
|
|
150
|
+
|
|
151
|
+
def _error_response(request: Request, code: ErrorCode, message: str, status: int, headers: dict[str, str] | None = None) -> JSONResponse:
|
|
152
|
+
request_id = getattr(request.state, "request_id", None) or uuid4().hex
|
|
153
|
+
state = _state(request)
|
|
154
|
+
state.metrics.record_error(code.value)
|
|
155
|
+
log_event("error", request_id=request_id, code=code.value, status=status)
|
|
156
|
+
body = {"error": ErrorBody(code=code, message=message, request_id=request_id).model_dump(mode="json")}
|
|
157
|
+
return JSONResponse(status_code=status, content=body, headers=headers)
|
|
158
|
+
|
|
159
|
+
|
|
160
|
+
def create_app(settings: Settings | None = None, scorer: Scorer | None = None) -> FastAPI:
|
|
161
|
+
"""Build the app. Pass ``scorer`` to skip loading the model from ``settings.model_dir``."""
|
|
162
|
+
settings = settings or Settings.from_env()
|
|
163
|
+
|
|
164
|
+
@asynccontextmanager
|
|
165
|
+
async def lifespan(app: FastAPI) -> AsyncIterator[None]:
|
|
166
|
+
state: _ServiceState = app.state.spamkit
|
|
167
|
+
if state.scorer is None and scorer is None:
|
|
168
|
+
try:
|
|
169
|
+
state.scorer = await run_in_threadpool(SpamtonClassifier, settings.model_dir)
|
|
170
|
+
log_event("model_loaded", model=state.scorer.model_id, model_version=state.scorer.model_version)
|
|
171
|
+
except Exception as exc: # noqa: BLE001 - any load failure must leave the service not-ready, not crash it
|
|
172
|
+
state.load_error = f"{type(exc).__name__}: {exc}"
|
|
173
|
+
logger.error(
|
|
174
|
+
json.dumps(
|
|
175
|
+
{
|
|
176
|
+
"event": "model_load_failed",
|
|
177
|
+
"model_dir": str(settings.model_dir),
|
|
178
|
+
"reason": state.load_error,
|
|
179
|
+
},
|
|
180
|
+
sort_keys=True,
|
|
181
|
+
)
|
|
182
|
+
)
|
|
183
|
+
yield
|
|
184
|
+
|
|
185
|
+
app = FastAPI(
|
|
186
|
+
title="Spamkit",
|
|
187
|
+
version=__version__,
|
|
188
|
+
summary="Spam scoring for user-submitted text.",
|
|
189
|
+
lifespan=lifespan,
|
|
190
|
+
)
|
|
191
|
+
app.state.spamkit = _ServiceState(settings)
|
|
192
|
+
app.state.spamkit.scorer = scorer
|
|
193
|
+
|
|
194
|
+
@app.middleware("http")
|
|
195
|
+
async def request_context(request: Request, call_next: Any) -> Any:
|
|
196
|
+
incoming = request.headers.get("x-request-id")
|
|
197
|
+
request_id = incoming if incoming and _REQUEST_ID.fullmatch(incoming) else uuid4().hex
|
|
198
|
+
request.state.request_id = request_id
|
|
199
|
+
started = time.perf_counter()
|
|
200
|
+
response = await call_next(request)
|
|
201
|
+
response.headers["X-Request-Id"] = request_id
|
|
202
|
+
route = getattr(request.scope.get("route"), "path", "unmatched")
|
|
203
|
+
_state(request).metrics.record_request(route, response.status_code)
|
|
204
|
+
log_event(
|
|
205
|
+
"request",
|
|
206
|
+
request_id=request_id,
|
|
207
|
+
method=request.method,
|
|
208
|
+
route=route,
|
|
209
|
+
status=response.status_code,
|
|
210
|
+
latency_ms=round((time.perf_counter() - started) * 1000, 1),
|
|
211
|
+
)
|
|
212
|
+
return response
|
|
213
|
+
|
|
214
|
+
@app.exception_handler(ApiError)
|
|
215
|
+
async def handle_api_error(request: Request, exc: ApiError) -> JSONResponse:
|
|
216
|
+
return _error_response(request, exc.code, exc.message, exc.status, exc.headers)
|
|
217
|
+
|
|
218
|
+
@app.exception_handler(Exception)
|
|
219
|
+
async def handle_unexpected(request: Request, exc: Exception) -> JSONResponse:
|
|
220
|
+
# Log only the exception type: its message could contain request content.
|
|
221
|
+
log_event("unhandled_exception", exception=type(exc).__name__)
|
|
222
|
+
return _error_response(request, ErrorCode.INTERNAL_ERROR, "An internal error occurred.", 500)
|
|
223
|
+
|
|
224
|
+
@app.get("/healthz", response_model=HealthResponse)
|
|
225
|
+
async def healthz() -> HealthResponse:
|
|
226
|
+
return HealthResponse(status="ok")
|
|
227
|
+
|
|
228
|
+
@app.get("/readyz", response_model=ReadyResponse, responses={503: {"model": ErrorBody}})
|
|
229
|
+
async def readyz(request: Request) -> Any:
|
|
230
|
+
state = _state(request)
|
|
231
|
+
if state.scorer is None:
|
|
232
|
+
raise ApiError(ErrorCode.MODEL_NOT_READY, "The model is not loaded. Check the startup logs.")
|
|
233
|
+
return ReadyResponse(status="ready", model=state.scorer.model_id, model_version=state.scorer.model_version)
|
|
234
|
+
|
|
235
|
+
@app.get("/v1/info", response_model=InfoResponse)
|
|
236
|
+
async def info(request: Request) -> InfoResponse:
|
|
237
|
+
_require_auth(request, settings)
|
|
238
|
+
state = _state(request)
|
|
239
|
+
thresholds = settings.thresholds
|
|
240
|
+
return InfoResponse(
|
|
241
|
+
api_version=API_VERSION,
|
|
242
|
+
model=MODEL_ID,
|
|
243
|
+
model_version=state.scorer.model_version if state.scorer else None,
|
|
244
|
+
max_tokens=state.scorer.max_tokens if state.scorer else 512,
|
|
245
|
+
max_text_bytes=settings.max_text_bytes,
|
|
246
|
+
max_batch_items=settings.max_batch_items,
|
|
247
|
+
thresholds={"clean": thresholds.clean, "spam": thresholds.spam},
|
|
248
|
+
threshold_config_version=thresholds.version,
|
|
249
|
+
service_version=__version__,
|
|
250
|
+
)
|
|
251
|
+
|
|
252
|
+
@app.post("/v1/score", response_model=ScoreResult)
|
|
253
|
+
async def score(request: Request) -> ScoreResult:
|
|
254
|
+
_require_auth(request, settings)
|
|
255
|
+
state = _state(request)
|
|
256
|
+
item = parse_score_input(await _read_json(request, settings.max_body_bytes), settings.max_text_bytes)
|
|
257
|
+
scorer, scored, latency_ms = await _score_texts(state, [item.text])
|
|
258
|
+
return _build_result(state, scorer, item, scored[0], latency_ms, request.state.request_id)
|
|
259
|
+
|
|
260
|
+
@app.post("/v1/score/batch", response_model=BatchResponse)
|
|
261
|
+
async def score_batch(request: Request) -> BatchResponse:
|
|
262
|
+
_require_auth(request, settings)
|
|
263
|
+
state = _state(request)
|
|
264
|
+
payload = await _read_json(request, settings.max_body_bytes)
|
|
265
|
+
items = payload.get("items") if isinstance(payload, dict) else None
|
|
266
|
+
if not isinstance(items, list) or not items:
|
|
267
|
+
raise ApiError(ErrorCode.INVALID_REQUEST, "Body must be a JSON object with a non-empty items array.")
|
|
268
|
+
if len(items) > settings.max_batch_items:
|
|
269
|
+
raise ApiError(ErrorCode.INPUT_TOO_LARGE, f"items may contain at most {settings.max_batch_items} entries.")
|
|
270
|
+
|
|
271
|
+
# Validate each item separately so one bad record does not discard the rest.
|
|
272
|
+
parsed: list[ScoreInput | ApiError] = []
|
|
273
|
+
for raw in items:
|
|
274
|
+
try:
|
|
275
|
+
parsed.append(parse_score_input(raw, settings.max_text_bytes))
|
|
276
|
+
except ApiError as exc:
|
|
277
|
+
parsed.append(exc)
|
|
278
|
+
|
|
279
|
+
valid = [(position, entry) for position, entry in enumerate(parsed) if isinstance(entry, ScoreInput)]
|
|
280
|
+
scored_by_position: dict[int, ScoredText] = {}
|
|
281
|
+
scorer: Scorer | None = None
|
|
282
|
+
latency_ms = 0.0
|
|
283
|
+
if valid:
|
|
284
|
+
scorer, scored, latency_ms = await _score_texts(state, [entry.text for _, entry in valid])
|
|
285
|
+
scored_by_position = {position: result for (position, _), result in zip(valid, scored)}
|
|
286
|
+
|
|
287
|
+
results: list[BatchItemOk | BatchItemError] = []
|
|
288
|
+
for position, entry in enumerate(parsed):
|
|
289
|
+
if isinstance(entry, ApiError):
|
|
290
|
+
results.append(
|
|
291
|
+
BatchItemError(
|
|
292
|
+
id=_caller_id(items[position]),
|
|
293
|
+
error=ErrorBody(code=entry.code, message=entry.message, request_id=request.state.request_id),
|
|
294
|
+
)
|
|
295
|
+
)
|
|
296
|
+
elif scorer is not None:
|
|
297
|
+
result = _build_result(
|
|
298
|
+
state, scorer, entry, scored_by_position[position], latency_ms, request.state.request_id
|
|
299
|
+
)
|
|
300
|
+
results.append(BatchItemOk(result=result))
|
|
301
|
+
return BatchResponse(results=results)
|
|
302
|
+
|
|
303
|
+
@app.get("/metrics", response_class=PlainTextResponse)
|
|
304
|
+
async def metrics(request: Request) -> str:
|
|
305
|
+
_require_auth(request, settings)
|
|
306
|
+
return _state(request).metrics.render()
|
|
307
|
+
|
|
308
|
+
return app
|
|
309
|
+
|
|
310
|
+
|
|
311
|
+
def _caller_id(raw: Any) -> str | None:
|
|
312
|
+
if isinstance(raw, dict) and isinstance(raw.get("id"), str) and len(raw["id"]) <= 256:
|
|
313
|
+
return raw["id"]
|
|
314
|
+
return None
|
spamkit/classifier.py
ADDED
|
@@ -0,0 +1,122 @@
|
|
|
1
|
+
"""Model runtime. The API depends only on the ``Scorer`` protocol, so the classifier can be swapped."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import json
|
|
6
|
+
import threading
|
|
7
|
+
from collections.abc import Sequence
|
|
8
|
+
from dataclasses import dataclass
|
|
9
|
+
from pathlib import Path
|
|
10
|
+
from typing import Protocol
|
|
11
|
+
|
|
12
|
+
MODEL_ID = "spamton"
|
|
13
|
+
MODEL_SOURCE = "tiagozip/spamton"
|
|
14
|
+
# Pinned revision of tiagozip/spamton (AGPL-3.0). Changing it is a deliberate, reviewed act.
|
|
15
|
+
MODEL_REVISION = "7aa94d049d24bbc302c5a69e63cafe6c44d49bc1"
|
|
16
|
+
MANIFEST_NAME = "spamkit-model.json"
|
|
17
|
+
|
|
18
|
+
# Preprocessing mirrors the model author's reference script: clip characters first, then
|
|
19
|
+
# truncate to 512 tokens. Index 1 of the softmax is the spam class.
|
|
20
|
+
MAX_CHARS = 6000
|
|
21
|
+
MAX_TOKENS = 512
|
|
22
|
+
INFERENCE_CHUNK = 16
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
class ModelLoadError(RuntimeError):
|
|
26
|
+
"""Raised when the model directory is missing, incomplete, or does not match the expected model."""
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
@dataclass(frozen=True)
|
|
30
|
+
class ScoredText:
|
|
31
|
+
score: float
|
|
32
|
+
truncated: bool
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
class Scorer(Protocol):
|
|
36
|
+
model_id: str
|
|
37
|
+
model_version: str
|
|
38
|
+
max_tokens: int
|
|
39
|
+
|
|
40
|
+
def score_many(self, texts: Sequence[str]) -> list[ScoredText]: ...
|
|
41
|
+
|
|
42
|
+
|
|
43
|
+
def read_manifest(model_dir: Path) -> dict[str, str]:
|
|
44
|
+
manifest_path = model_dir / MANIFEST_NAME
|
|
45
|
+
if not manifest_path.is_file():
|
|
46
|
+
raise ModelLoadError(
|
|
47
|
+
f"No {MANIFEST_NAME} in {model_dir}. Run `spamkit fetch-model --dest {model_dir}` first."
|
|
48
|
+
)
|
|
49
|
+
try:
|
|
50
|
+
manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
|
|
51
|
+
except (OSError, ValueError) as exc:
|
|
52
|
+
raise ModelLoadError(f"Could not read {manifest_path}: {exc}") from exc
|
|
53
|
+
if manifest.get("model") != MODEL_ID or not manifest.get("revision"):
|
|
54
|
+
raise ModelLoadError(f"{manifest_path} does not describe the {MODEL_ID} model.")
|
|
55
|
+
return manifest
|
|
56
|
+
|
|
57
|
+
|
|
58
|
+
class SpamtonClassifier:
|
|
59
|
+
"""Runs spamton on CPU from a local directory. Never downloads anything at runtime."""
|
|
60
|
+
|
|
61
|
+
model_id = MODEL_ID
|
|
62
|
+
max_tokens = MAX_TOKENS
|
|
63
|
+
|
|
64
|
+
def __init__(self, model_dir: Path) -> None:
|
|
65
|
+
manifest = read_manifest(model_dir)
|
|
66
|
+
self.model_version: str = manifest["revision"]
|
|
67
|
+
|
|
68
|
+
import torch
|
|
69
|
+
from transformers import AutoModelForSequenceClassification, AutoTokenizer
|
|
70
|
+
|
|
71
|
+
self._torch = torch
|
|
72
|
+
self._tokenizer = AutoTokenizer.from_pretrained(model_dir, local_files_only=True)
|
|
73
|
+
self._model = AutoModelForSequenceClassification.from_pretrained(
|
|
74
|
+
model_dir, local_files_only=True
|
|
75
|
+
).eval()
|
|
76
|
+
# One forward pass at a time: CPU inference does not gain from concurrent calls.
|
|
77
|
+
self._lock = threading.Lock()
|
|
78
|
+
# The first forward pass pays one-time allocation cost; do it before readiness flips.
|
|
79
|
+
self.score_many(["warm-up"])
|
|
80
|
+
|
|
81
|
+
def score_many(self, texts: Sequence[str]) -> list[ScoredText]:
|
|
82
|
+
results: list[ScoredText] = []
|
|
83
|
+
for start in range(0, len(texts), INFERENCE_CHUNK):
|
|
84
|
+
results.extend(self._score_chunk(texts[start : start + INFERENCE_CHUNK]))
|
|
85
|
+
return results
|
|
86
|
+
|
|
87
|
+
def _score_chunk(self, texts: Sequence[str]) -> list[ScoredText]:
|
|
88
|
+
clipped = [text[:MAX_CHARS] for text in texts]
|
|
89
|
+
# The Rust tokenizer is not safe to call from several threads at once ("Already borrowed"),
|
|
90
|
+
# so tokenization and the forward pass both run under the same lock.
|
|
91
|
+
with self._lock:
|
|
92
|
+
truncated = [len(self._tokenizer(text)["input_ids"]) > MAX_TOKENS for text in clipped]
|
|
93
|
+
encoded = self._tokenizer(
|
|
94
|
+
clipped,
|
|
95
|
+
truncation=True,
|
|
96
|
+
max_length=MAX_TOKENS,
|
|
97
|
+
padding=True,
|
|
98
|
+
return_tensors="pt",
|
|
99
|
+
)
|
|
100
|
+
with self._torch.inference_mode():
|
|
101
|
+
spam_probabilities = (
|
|
102
|
+
self._model(**encoded).logits.float().softmax(-1)[:, 1].tolist()
|
|
103
|
+
)
|
|
104
|
+
return [
|
|
105
|
+
ScoredText(score=float(score), truncated=was_truncated)
|
|
106
|
+
for score, was_truncated in zip(spam_probabilities, truncated)
|
|
107
|
+
]
|
|
108
|
+
|
|
109
|
+
|
|
110
|
+
def fetch_model(dest: Path, revision: str = MODEL_REVISION) -> None:
|
|
111
|
+
"""Download spamton at a pinned revision and write the manifest the server requires."""
|
|
112
|
+
from huggingface_hub import snapshot_download
|
|
113
|
+
|
|
114
|
+
dest.mkdir(parents=True, exist_ok=True)
|
|
115
|
+
snapshot_download(repo_id=MODEL_SOURCE, revision=revision, local_dir=dest)
|
|
116
|
+
manifest = {
|
|
117
|
+
"model": MODEL_ID,
|
|
118
|
+
"source": MODEL_SOURCE,
|
|
119
|
+
"revision": revision,
|
|
120
|
+
"license": "AGPL-3.0 (see LICENSE and NOTICE in this directory)",
|
|
121
|
+
}
|
|
122
|
+
(dest / MANIFEST_NAME).write_text(json.dumps(manifest, indent=2) + "\n", encoding="utf-8")
|
spamkit/cli.py
ADDED
|
@@ -0,0 +1,199 @@
|
|
|
1
|
+
"""Command-line entry point: ``spamkit serve | fetch-model | evaluate``."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import argparse
|
|
6
|
+
import json
|
|
7
|
+
import logging
|
|
8
|
+
import sys
|
|
9
|
+
from dataclasses import replace
|
|
10
|
+
from pathlib import Path
|
|
11
|
+
from typing import Any
|
|
12
|
+
|
|
13
|
+
from spamkit import __version__
|
|
14
|
+
from spamkit.evaluation import (
|
|
15
|
+
DEFAULT_SWEEP,
|
|
16
|
+
EvaluationInputError,
|
|
17
|
+
build_report,
|
|
18
|
+
load_samples,
|
|
19
|
+
)
|
|
20
|
+
from spamkit.settings import Settings
|
|
21
|
+
|
|
22
|
+
EXIT_USAGE = 2
|
|
23
|
+
EXIT_FAILURE = 1
|
|
24
|
+
SERVER_EXTRA_HINT = "This command needs the model runtime. Install with: pip install 'spamkit[server]'"
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
def main(argv: list[str] | None = None) -> int:
|
|
28
|
+
parser = _build_parser()
|
|
29
|
+
args = parser.parse_args(argv)
|
|
30
|
+
try:
|
|
31
|
+
return args.handler(args)
|
|
32
|
+
except ModuleNotFoundError as exc:
|
|
33
|
+
print(f"{SERVER_EXTRA_HINT} (missing module: {exc.name})", file=sys.stderr)
|
|
34
|
+
return EXIT_FAILURE
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
def _build_parser() -> argparse.ArgumentParser:
|
|
38
|
+
parser = argparse.ArgumentParser(prog="spamkit", description="Spam scoring for user-submitted text.")
|
|
39
|
+
parser.add_argument("--version", action="version", version=f"spamkit {__version__}")
|
|
40
|
+
commands = parser.add_subparsers(dest="command", required=True)
|
|
41
|
+
|
|
42
|
+
serve = commands.add_parser("serve", help="Run the scoring API.")
|
|
43
|
+
serve.add_argument("--host", help="Bind address (default: SPAMKIT_HOST or 127.0.0.1).")
|
|
44
|
+
serve.add_argument("--port", type=int, help="Port (default: SPAMKIT_PORT or 8080).")
|
|
45
|
+
serve.add_argument("--model-dir", type=Path, help="Model directory (default: SPAMKIT_MODEL_DIR).")
|
|
46
|
+
serve.set_defaults(handler=_serve)
|
|
47
|
+
|
|
48
|
+
fetch = commands.add_parser("fetch-model", help="Download spamton at the pinned revision.")
|
|
49
|
+
fetch.add_argument("--dest", type=Path, required=True, help="Directory to write the model into.")
|
|
50
|
+
fetch.add_argument("--revision", help="Override the pinned revision (not recommended).")
|
|
51
|
+
fetch.set_defaults(handler=_fetch_model)
|
|
52
|
+
|
|
53
|
+
evaluate = commands.add_parser("evaluate", help="Score a local labeled file and report metrics.")
|
|
54
|
+
evaluate.add_argument("path", type=Path, help="CSV or JSONL file with text and label columns. Row numbers in output are file line numbers.")
|
|
55
|
+
evaluate.add_argument("--threshold", type=float, default=0.5, help="Predict spam when score >= threshold.")
|
|
56
|
+
evaluate.add_argument(
|
|
57
|
+
"--sweep",
|
|
58
|
+
type=float,
|
|
59
|
+
nargs="+",
|
|
60
|
+
default=list(DEFAULT_SWEEP),
|
|
61
|
+
help="Thresholds for the trade-off table (default: 0.05 to 0.95).",
|
|
62
|
+
)
|
|
63
|
+
evaluate.add_argument("--model-dir", type=Path, help="Model directory (default: SPAMKIT_MODEL_DIR).")
|
|
64
|
+
evaluate.add_argument("--output", type=Path, help="Write the full report as JSON to this path.")
|
|
65
|
+
evaluate.add_argument(
|
|
66
|
+
"--show-text",
|
|
67
|
+
action="store_true",
|
|
68
|
+
help="Print the text of false positives and negatives. Off by default.",
|
|
69
|
+
)
|
|
70
|
+
evaluate.set_defaults(handler=_evaluate)
|
|
71
|
+
return parser
|
|
72
|
+
|
|
73
|
+
|
|
74
|
+
def _serve(args: argparse.Namespace) -> int:
|
|
75
|
+
import uvicorn
|
|
76
|
+
|
|
77
|
+
from spamkit.app import create_app
|
|
78
|
+
|
|
79
|
+
settings = Settings.from_env()
|
|
80
|
+
overrides: dict[str, Any] = {}
|
|
81
|
+
if args.host:
|
|
82
|
+
overrides["host"] = args.host
|
|
83
|
+
if args.port:
|
|
84
|
+
overrides["port"] = args.port
|
|
85
|
+
if args.model_dir:
|
|
86
|
+
overrides["model_dir"] = args.model_dir
|
|
87
|
+
settings = replace(settings, **overrides)
|
|
88
|
+
|
|
89
|
+
logging.basicConfig(level=logging.INFO, format="%(message)s")
|
|
90
|
+
app = create_app(settings)
|
|
91
|
+
# One worker only: each worker would load its own copy of the model.
|
|
92
|
+
uvicorn.run(
|
|
93
|
+
app,
|
|
94
|
+
host=settings.host,
|
|
95
|
+
port=settings.port,
|
|
96
|
+
workers=1,
|
|
97
|
+
log_level="warning",
|
|
98
|
+
access_log=False,
|
|
99
|
+
timeout_graceful_shutdown=30,
|
|
100
|
+
)
|
|
101
|
+
return 0
|
|
102
|
+
|
|
103
|
+
|
|
104
|
+
def _fetch_model(args: argparse.Namespace) -> int:
|
|
105
|
+
from spamkit.classifier import MODEL_REVISION, fetch_model
|
|
106
|
+
|
|
107
|
+
revision = args.revision or MODEL_REVISION
|
|
108
|
+
print(f"Fetching spamton @ {revision} into {args.dest} ...")
|
|
109
|
+
fetch_model(args.dest, revision=revision)
|
|
110
|
+
print(f"Done. Manifest written to {args.dest / 'spamkit-model.json'}")
|
|
111
|
+
return 0
|
|
112
|
+
|
|
113
|
+
|
|
114
|
+
def _evaluate(args: argparse.Namespace) -> int:
|
|
115
|
+
from spamkit.classifier import SpamtonClassifier
|
|
116
|
+
|
|
117
|
+
if not 0.0 <= args.threshold <= 1.0:
|
|
118
|
+
print("--threshold must be between 0 and 1.", file=sys.stderr)
|
|
119
|
+
return EXIT_USAGE
|
|
120
|
+
try:
|
|
121
|
+
samples = load_samples(args.path)
|
|
122
|
+
except EvaluationInputError as exc:
|
|
123
|
+
print(f"Input error: {exc}", file=sys.stderr)
|
|
124
|
+
return EXIT_USAGE
|
|
125
|
+
|
|
126
|
+
settings = Settings.from_env()
|
|
127
|
+
model_dir = args.model_dir or settings.model_dir
|
|
128
|
+
try:
|
|
129
|
+
scorer = SpamtonClassifier(model_dir)
|
|
130
|
+
except Exception as exc: # noqa: BLE001 - any load failure becomes a clear message and exit code
|
|
131
|
+
print(f"Could not load the model from {model_dir}: {exc}", file=sys.stderr)
|
|
132
|
+
return EXIT_FAILURE
|
|
133
|
+
|
|
134
|
+
scored = scorer.score_many([sample.text for sample in samples])
|
|
135
|
+
scores = [item.score for item in scored]
|
|
136
|
+
report = build_report(
|
|
137
|
+
samples,
|
|
138
|
+
scores,
|
|
139
|
+
threshold=args.threshold,
|
|
140
|
+
sweep=sorted(set(args.sweep)),
|
|
141
|
+
model=scorer.model_id,
|
|
142
|
+
model_version=scorer.model_version,
|
|
143
|
+
software_version=__version__,
|
|
144
|
+
show_text=args.show_text,
|
|
145
|
+
)
|
|
146
|
+
|
|
147
|
+
_print_report(report, show_text=args.show_text)
|
|
148
|
+
if args.output:
|
|
149
|
+
args.output.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
|
|
150
|
+
print(f"\nFull report written to {args.output}")
|
|
151
|
+
return 0
|
|
152
|
+
|
|
153
|
+
|
|
154
|
+
def _fmt(value: float | None) -> str:
|
|
155
|
+
return "n/a" if value is None else f"{value:.3f}"
|
|
156
|
+
|
|
157
|
+
|
|
158
|
+
def _print_report(report: dict[str, Any], *, show_text: bool) -> None:
|
|
159
|
+
dataset = report["dataset"]
|
|
160
|
+
matrix = report["confusion"]
|
|
161
|
+
metrics = report["metrics"]
|
|
162
|
+
print("Spamkit evaluation")
|
|
163
|
+
print(f"Dataset: {dataset['count']} rows (spam {dataset['spam']}, clean {dataset['clean']})")
|
|
164
|
+
print(f"Model: {report['model']['id']} @ {report['model']['version']} software {report['software_version']}")
|
|
165
|
+
print(f"Threshold: {report['threshold']:.2f} (predict spam when score >= threshold)")
|
|
166
|
+
print()
|
|
167
|
+
print(f"Confusion matrix TP {matrix['tp']} FP {matrix['fp']} TN {matrix['tn']} FN {matrix['fn']}")
|
|
168
|
+
print(
|
|
169
|
+
f"Precision {_fmt(metrics['precision'])} Recall {_fmt(metrics['recall'])} "
|
|
170
|
+
f"Specificity {_fmt(metrics['specificity'])} F1 {_fmt(metrics['f1'])}"
|
|
171
|
+
)
|
|
172
|
+
print(
|
|
173
|
+
f"False-positive rate {_fmt(metrics['false_positive_rate'])} "
|
|
174
|
+
f"False-negative rate {_fmt(metrics['false_negative_rate'])}"
|
|
175
|
+
)
|
|
176
|
+
print()
|
|
177
|
+
print(f"{'threshold':>9} {'precision':>9} {'recall':>6} {'FPR':>6} {'FNR':>6} {'F1':>6}")
|
|
178
|
+
for point in report["sweep"]:
|
|
179
|
+
print(
|
|
180
|
+
f"{point['threshold']:>9.2f} {_fmt(point['precision']):>9} {_fmt(point['recall']):>6} "
|
|
181
|
+
f"{_fmt(point['false_positive_rate']):>6} {_fmt(point['false_negative_rate']):>6} {_fmt(point['f1']):>6}"
|
|
182
|
+
)
|
|
183
|
+
print()
|
|
184
|
+
_print_errors("False positives (clean flagged as spam)", report["false_positives"], show_text)
|
|
185
|
+
_print_errors("False negatives (spam not flagged)", report["false_negatives"], show_text)
|
|
186
|
+
print(dataset["note"])
|
|
187
|
+
|
|
188
|
+
|
|
189
|
+
def _print_errors(title: str, rows: list[dict[str, Any]], show_text: bool) -> None:
|
|
190
|
+
print(f"{title}: {len(rows)}")
|
|
191
|
+
for row in rows:
|
|
192
|
+
label = f"row {row['row']}" + (f" (id {row['id']})" if row["id"] else "")
|
|
193
|
+
print(f" {label} score {row['score']:.3f}")
|
|
194
|
+
if show_text:
|
|
195
|
+
print(f" text: {row['text']!r}")
|
|
196
|
+
|
|
197
|
+
|
|
198
|
+
if __name__ == "__main__":
|
|
199
|
+
raise SystemExit(main())
|