| import json |
| import sys |
|
|
| import numpy as np |
| import onnxruntime as ort |
| import sentencepiece as spm |
| from huggingface_hub import hf_hub_download |
|
|
| REPO = "Prosperitytech/lighthouse-keeper" |
|
|
| cfg = json.load(open(hf_hub_download(REPO, "config.json"))) |
| sp = spm.SentencePieceProcessor(model_file=hf_hub_download(REPO, "keeper.spm")) |
| session = ort.InferenceSession(hf_hub_download(REPO, "keeper.onnx"), providers=["CPUExecutionProvider"]) |
|
|
|
|
| def score(text): |
| ids = [cfg["bos_id"]] + sp.encode(text)[: cfg["max_tokens"] - 1] |
| probs = session.run(None, {"input_ids": np.array([ids], dtype=np.int64)})[0][0] |
| return {cat: dict(zip(cfg["levels"], map(float, p))) for cat, p in zip(cfg["categories"], probs)} |
|
|
|
|
| def flag(text, min_risk=cfg["suggested_min_risk"]): |
| priority = cfg["tie_break_order"] |
| best = None |
| for cat, p in score(text).items(): |
| level = max(p, key=p.get) |
| risk = 1 - p["none"] |
| if level != "none" and risk >= min_risk: |
| key = (cfg["levels"].index(level), p[level] + 0.06 * (1 - priority.index(cat) / len(priority))) |
| if best is None or key > best[1]: |
| best = (cat, key, p[level]) |
| return None if best is None else {"category": best[0], "level": cfg["levels"][best[1][0]], "confidence": round(best[2], 3)} |
|
|
|
|
| if __name__ == "__main__": |
| print(flag(" ".join(sys.argv[1:]) or "send me your BVN make I help you collect the money")) |
|
|
