@swapai/core 0.1.0 → 0.2.1

package/dist/{types-DTbU1MEi.d.ts → types-DkvGbscw.d.ts}

@@ -24,6 +24,7 @@ type ReferenceClassifier<Result extends ResultValue> = (input: string) => Result
 interface Classifier<Result extends ResultValue> {
     isTrained(): boolean;
     logClassification(input: string, result: Result): void;
+    clearTrainingData(): void;
     classify(input: string, referenceClassifier?: ReferenceClassifier<Result>): Promise<Result>;
     flush(): Promise<void>;
     close(): Promise<void>;

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@swapai/core",
-  "version": "0.1.0",
+  "version": "0.2.1",
   "description": "Replace paid classifiers with locally trained Needle 2 classifiers.",
   "type": "module",
   "license": "MIT",

package/python/swapai_worker.py

@@ -3,9 +3,12 @@ from __future__ import annotations
 
 import argparse
 import json
+import math
 import os
+import sqlite3
 import sys
 import types
+from contextlib import contextmanager
 from pathlib import Path
 from typing import Any
 
@@ -28,52 +31,106 @@ def require_version(needle: Any) -> None:
         )
 
 
-def train(args: argparse.Namespace) -> None:
-    import needle
-    from needle.model.finetune import build_main, finetune_local
+@contextmanager
+def artifact_write_lock(args: argparse.Namespace):
+    lock_database = sqlite3.connect(args.artifact_lock_database, timeout=30)
+    main_database: sqlite3.Connection | None = None
+    try:
+        lock_database.execute("BEGIN IMMEDIATE")
+        main_database = sqlite3.connect(f"file:{args.main_database}?mode=ro", uri=True)
+        row = main_database.execute(
+            "SELECT data_epoch, clear_pending FROM classifiers WHERE name = ?",
+            (args.classifier_name,),
+        ).fetchone()
+        if row is None or row[0] != args.expected_epoch or row[1] != 0:
+            raise RuntimeError("Classifier training was superseded by data erasure")
+        yield
+    finally:
+        if main_database is not None:
+            main_database.close()
+        try:
+            lock_database.rollback()
+        finally:
+            lock_database.close()
 
-    require_version(needle)
-    checkpoint_dir = Path(args.checkpoint_dir).resolve()
-    checkpoint_dir.mkdir(parents=True, exist_ok=True)
-    checkpoint = checkpoint_dir / "needle2.pkl"
-    output = Path(args.output).resolve()
-    output.parent.mkdir(parents=True, exist_ok=True)
-    adapter = output.parent / "swapai-lora.pkl"
-    with open(args.training_data, "r", encoding="utf-8") as handle:
-        example_count = sum(1 for line in handle if line.strip())
-    batch_size = min(16, max(1, (example_count + 9) // 10))
-
-    finetune_local(
-        types.SimpleNamespace(
-            jsonl_path=str(Path(args.training_data).resolve()),
-            checkpoint=str(checkpoint),
-            epochs=args.epochs,
-            batch_size=batch_size,
-            lr=1e-4,
-            lora_rank=16,
-            lora_alpha=32.0,
-            max_len=1024,
-            val_split=0.0,
-            seed=0,
-            generate=0,
-            model="deepseek/deepseek-v4-flash",
-            workers=1,
-            checkpoint_dir=str(checkpoint_dir),
-            out=str(adapter),
-            qat_bits="auto",
+
+def train(args: argparse.Namespace) -> None:
+    with artifact_write_lock(args):
+        import needle
+        from needle.model.finetune import build_main, finetune_local
+
+        require_version(needle)
+        checkpoint_dir = Path(args.checkpoint_dir).resolve()
+        checkpoint_dir.mkdir(parents=True, exist_ok=True)
+        checkpoint = checkpoint_dir / "needle2.pkl"
+        output = Path(args.output).resolve()
+        output.parent.mkdir(parents=True, exist_ok=True)
+        training_input = sys.stdin.read()
+        number_labels: dict[str, Any] | None = None
+        if args.numeric_labels_from_stdin:
+            labels_json, separator, training_input = training_input.partition("\n")
+            if not separator:
+                raise RuntimeError("Numeric label map is missing")
+            labels = json.loads(labels_json)
+            values = labels.get("values") if isinstance(labels, dict) else None
+            if (
+                not isinstance(labels, dict)
+                or labels.get("format") != 1
+                or not isinstance(values, list)
+                or not values
+                or any(
+                    not isinstance(value, (int, float))
+                    or isinstance(value, bool)
+                    or not math.isfinite(value)
+                    for value in values
+                )
+                or values != sorted(set(values))
+            ):
+                raise RuntimeError("Numeric label map has an invalid format")
+            number_labels = labels
+        training_data = Path(args.training_data).resolve()
+        training_data.write_text(training_input, encoding="utf-8")
+        adapter = output.parent / "swapai-lora.pkl"
+        with training_data.open("r", encoding="utf-8") as handle:
+            example_count = sum(1 for line in handle if line.strip())
+        batch_size = min(16, max(1, (example_count + 9) // 10))
+
+        finetune_local(
+            types.SimpleNamespace(
+                jsonl_path=str(training_data),
+                checkpoint=str(checkpoint),
+                epochs=args.epochs,
+                batch_size=batch_size,
+                lr=1e-4,
+                lora_rank=16,
+                lora_alpha=32.0,
+                max_len=1024,
+                val_split=0.0,
+                seed=0,
+                generate=0,
+                model="deepseek/deepseek-v4-flash",
+                workers=1,
+                checkpoint_dir=str(checkpoint_dir),
+                out=str(adapter),
+                qat_bits="auto",
+            )
         )
-    )
-    build_main(
-        types.SimpleNamespace(
-            checkpoint=str(checkpoint),
-            lora=str(adapter),
-            out=str(output),
-            upload=False,
-            bits=None,
+        build_main(
+            types.SimpleNamespace(
+                checkpoint=str(checkpoint),
+                lora=str(adapter),
+                out=str(output),
+                upload=False,
+                bits=None,
+            )
         )
-    )
-    if not output.is_file():
-        raise RuntimeError("Needle did not create the requested .cact model")
+        if not output.is_file():
+            raise RuntimeError("Needle did not create the requested .cact model")
+        if number_labels is not None:
+            Path(f"{output}.numbers.json").write_text(
+                json.dumps(number_labels, separators=(",", ":")),
+                encoding="utf-8",
+            )
 
 
 def classification_result(response: Any) -> Any:
@@ -147,6 +204,11 @@ def parser() -> argparse.ArgumentParser:
     train_parser.add_argument("--output", required=True)
     train_parser.add_argument("--checkpoint-dir", required=True)
     train_parser.add_argument("--epochs", type=int, default=10)
+    train_parser.add_argument("--artifact-lock-database", required=True)
+    train_parser.add_argument("--main-database", required=True)
+    train_parser.add_argument("--classifier-name", required=True)
+    train_parser.add_argument("--expected-epoch", type=int, required=True)
+    train_parser.add_argument("--numeric-labels-from-stdin", action="store_true")
 
     serve_parser = commands.add_parser("serve")
     serve_parser.add_argument("--model", required=True)