Text Classification
ONNX
English
multilingual
pico-type
pico_type
byte-level
content-classification
multi-head
classifier
edge-ai
code-detection
language-detection
clipboard
tiny
open-source
onnx-runtime
content-moderation
risk-detection
secrets-detection
prompt-injection-detection
pii-detection
pico-type-arth
Eval Results (legacy)
| """pico-type CLI: classify content from stdin/file/clipboard.""" | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import os | |
| import sys | |
| import numpy as np | |
| from .labels import ( | |
| ALL_HEADS, | |
| COARSE_LABELS, | |
| CODE_LANG_LABELS, | |
| FILE_MIME_LABELS, | |
| MODALITY_LABELS, | |
| RISK_LABELS, | |
| SUBTYPE_LABELS, | |
| TEXT_LANG_LABELS, | |
| ) | |
| LABEL_TABLES = { | |
| "coarse": COARSE_LABELS, | |
| "modality": MODALITY_LABELS, | |
| "subtype": SUBTYPE_LABELS, | |
| "code_lang": CODE_LANG_LABELS, | |
| "text_lang": TEXT_LANG_LABELS, | |
| "file_mime": FILE_MIME_LABELS, | |
| "risk": RISK_LABELS, | |
| } | |
| def load_onnx_model(tier: str = "base", model_dir: str = "checkpoints"): | |
| import onnxruntime as ort | |
| path = os.path.join(model_dir, f"picotype_{tier}.onnx") | |
| if not os.path.exists(path): | |
| raise FileNotFoundError(f"ONNX model not found: {path}") | |
| session = ort.InferenceSession(path) | |
| return session | |
| def load_torch_model(tier: str = "base", checkpoint: str = ""): | |
| import torch | |
| from .arch import PicoType, PicoTypeConfig | |
| cfg = PicoTypeConfig(max_bytes=1024) | |
| model = PicoType(cfg) | |
| ckpt = torch.load(checkpoint, map_location="cpu") | |
| model.load_state_dict(ckpt.get("model_state_dict", ckpt)) | |
| model.eval() | |
| return model, tier | |
| def run_onnx(session, data: bytes, max_bytes: int = 1024) -> dict: | |
| ids = np.frombuffer(data[:max_bytes], dtype=np.uint8).astype(np.int64) | |
| seq_len = len(ids) | |
| if seq_len > max_bytes: | |
| ids = ids[:max_bytes] | |
| seq_len = max_bytes | |
| padded = np.zeros(max_bytes, dtype=np.int64) | |
| padded[:seq_len] = ids | |
| mask = np.zeros(max_bytes, dtype=np.bool_) | |
| mask[:seq_len] = True | |
| feed = { | |
| "input_ids": padded[None, :], | |
| "attention_mask": mask[None, :], | |
| } | |
| outs = session.run(None, feed) | |
| result = {} | |
| for name, logits in zip(ALL_HEADS, outs): | |
| probs = _softmax(logits[0]) | |
| if name == "risk": | |
| result[name] = {LABEL_TABLES[name][i]: float(probs[i]) for i in range(len(probs))} | |
| else: | |
| idx = int(np.argmax(probs)) | |
| label = LABEL_TABLES[name][idx] | |
| result[name] = {"label": label, "confidence": float(probs[idx]), "index": idx} | |
| return result | |
| def _softmax(x): | |
| e = np.exp(x - np.max(x)) | |
| return e / e.sum() | |
| def run_torch(model, tier: str, data: bytes, max_bytes: int = 1024) -> dict: | |
| import torch | |
| model = model[0] if isinstance(model, tuple) else model | |
| ids = torch.tensor([list(data[:max_bytes])], dtype=torch.long) | |
| mask = torch.ones(1, ids.shape[1], dtype=torch.bool) | |
| with torch.no_grad(): | |
| logits_dict = model(ids, mask) | |
| out = {} | |
| for head in ALL_HEADS: | |
| logits = logits_dict[head] | |
| tier_logits = logits[tier] if isinstance(logits, dict) else logits | |
| if isinstance(tier_logits, dict): | |
| tier_logits = tier_logits[tier] | |
| probs = torch.softmax(tier_logits[0], dim=-1) | |
| if head == "risk": | |
| out[head] = {LABEL_TABLES[head][i]: float(probs[i]) for i in range(len(probs))} | |
| else: | |
| idx = int(torch.argmax(probs).item()) | |
| out[head] = {"label": LABEL_TABLES[head][idx], "confidence": float(probs[idx]), "index": idx} | |
| return out | |
| def read_input(args) -> bytes: | |
| if args.text: | |
| return args.text.encode("utf-8") | |
| if args.file: | |
| with open(args.file, "rb") as f: | |
| return f.read() | |
| if args.clip: | |
| import subprocess | |
| return subprocess.check_output(["pbpaste"]).rstrip(b"\n") | |
| if not sys.stdin.isatty(): | |
| return sys.stdin.buffer.read() | |
| raise ValueError("No input provided. Use --text, --file, --clip, or pipe content.") | |
| def build_parser(): | |
| p = argparse.ArgumentParser(prog="picotype", description="Classify content type and risk") | |
| p.add_argument("--text", "-t", help="Text string to classify") | |
| p.add_argument("--file", "-f", help="File path to classify") | |
| p.add_argument("--clip", "-c", action="store_true", help="Classify clipboard content") | |
| p.add_argument("--tier", default="base", choices=["tiny", "small", "base", "pro"], help="Model tier") | |
| p.add_argument("--model-dir", default="checkpoints", help="Directory with ONNX models") | |
| p.add_argument("--checkpoint", help="PyTorch checkpoint (fallback if no ONNX)") | |
| p.add_argument("--pretty", "-p", action="store_true", help="Pretty-print JSON output") | |
| return p | |
| def main(): | |
| args = build_parser().parse_args() | |
| try: | |
| data = read_input(args) | |
| except ValueError as e: | |
| print(e, file=sys.stderr) | |
| sys.exit(1) | |
| if not data.strip(): | |
| print('{"error": "empty input"}') | |
| sys.exit(0) | |
| onnx_path = os.path.join(args.model_dir, f"picotype_{args.tier}.onnx") | |
| if os.path.exists(onnx_path): | |
| session = load_onnx_model(args.tier, args.model_dir) | |
| result = run_onnx(session, data) | |
| elif args.checkpoint: | |
| model = load_torch_model(args.tier, args.checkpoint) | |
| result = run_torch(model, args.tier, data) | |
| else: | |
| print(f"ONNX model not found at {onnx_path}. Use --checkpoint to use PyTorch.", file=sys.stderr) | |
| sys.exit(1) | |
| result["text_length"] = len(data) | |
| result["tier"] = args.tier | |
| indent = 2 if args.pretty else None | |
| json.dump(result, sys.stdout, indent=indent, ensure_ascii=False) | |
| print() | |
| if __name__ == "__main__": | |
| main() | |