gheim 0.1.4 → 0.1.6

package/dist/cjs/index.cjs

@@ -58,7 +58,8 @@ __export(exports_src, {
   LocalDetector: () => LocalDetector,
   LABEL_TO_TAG: () => LABEL_TO_TAG,
   DEFAULT_BASE_URL: () => DEFAULT_BASE_URL,
-  CompositeDetector: () => CompositeDetector
+  CompositeDetector: () => CompositeDetector,
+  CalibratedDetector: () => CalibratedDetector
 });
 module.exports = __toCommonJS(exports_src);
 
@@ -156,8 +157,8 @@ class Session {
     }
     out.push(text.slice(cursor));
     let redacted = out.join("");
-    redacted = recoverDuplicateOccurrences(redacted, this.mapping);
-    return { redacted, merged };
+    const recovered = recoverDuplicateOccurrences(text, redacted, merged, this.mapping);
+    return { redacted: recovered.redacted, merged: recovered.merged };
   }
   restore(text) {
     if (Object.keys(this.mapping).length === 0)
@@ -216,14 +217,36 @@ function mergeAdjacent(text, spans) {
 function escapeRegex(s) {
   return s.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
 }
-function recoverDuplicateOccurrences(redacted, mapping) {
+function recoverDuplicateOccurrences(text, redacted, merged, mapping) {
+  const surfaceToLabel = new Map;
+  for (const m of merged)
+    surfaceToLabel.set(text.slice(m.start, m.end), m.label);
+  const covered = merged.map((m) => [m.start, m.end]);
+  const overlaps = (s, e) => covered.some(([cs, ce]) => s < ce && e > cs);
   const entries = Object.entries(mapping).filter(([, surface]) => surface && surface.length > 0).sort((a, b) => b[1].length - a[1].length);
   let result = redacted;
+  const additional = [];
   for (const [sentinel, surface] of entries) {
-    const pattern = new RegExp(`(?<![\\p{L}\\p{N}_])${escapeRegex(surface)}(?![\\p{L}\\p{N}_])`, "gu");
-    result = result.replace(pattern, sentinel);
+    const boundaryPattern = `(?<![\\p{L}\\p{N}_])${escapeRegex(surface)}(?![\\p{L}\\p{N}_])`;
+    result = result.replace(new RegExp(boundaryPattern, "gu"), sentinel);
+    const label = surfaceToLabel.get(surface);
+    if (!label)
+      continue;
+    const origRe = new RegExp(boundaryPattern, "gu");
+    let match;
+    while ((match = origRe.exec(text)) !== null) {
+      const start = match.index;
+      const end = start + surface.length;
+      if (overlaps(start, end))
+        continue;
+      additional.push({ start, end, label });
+      covered.push([start, end]);
+    }
   }
-  return result;
+  if (additional.length === 0)
+    return { redacted: result, merged };
+  const augmented = [...merged, ...additional].sort((a, b) => a.start - b.start);
+  return { redacted: result, merged: augmented };
 }
 // src/core/stream.ts
 var SENTINEL_AT = /<([A-Z][A-Z0-9_]*)_(\d+)>/y;
@@ -842,10 +865,281 @@ class CompositeDetector {
     return out;
   }
 }
+// src/detectors/calibrated.ts
+var DEFAULT_MODEL3 = "joelbarmettler/gheim-ch-560m";
+var DEFAULT_CHUNK_TOKENS2 = 510;
+var DEFAULT_OVERLAP_TOKENS2 = 32;
+function trimSpan2(text, start, end) {
+  while (start < end && /\s/.test(text[start] ?? ""))
+    start++;
+  while (end > start && /\s/.test(text[end - 1] ?? ""))
+    end--;
+  return [start, end];
+}
+
+class CalibratedDetector {
+  model;
+  oBias;
+  device;
+  dtype;
+  chunkTokens;
+  chunkOverlapTokens;
+  trimWhitespace;
+  actualDevice = null;
+  _model = null;
+  _tokenizer = null;
+  _id2label = {};
+  _oId = 0;
+  _loadPromise = null;
+  _onProgress;
+  constructor(opts = {}) {
+    this.model = opts.model ?? DEFAULT_MODEL3;
+    this.oBias = opts.oBias ?? 0.5;
+    this.device = opts.device ?? "auto";
+    this.dtype = opts.dtype ?? "fp32";
+    this.chunkTokens = opts.chunkTokens ?? DEFAULT_CHUNK_TOKENS2;
+    this.chunkOverlapTokens = opts.chunkOverlapTokens ?? DEFAULT_OVERLAP_TOKENS2;
+    this.trimWhitespace = opts.trimWhitespace ?? true;
+    this._onProgress = opts.onProgress;
+    if (this.chunkOverlapTokens >= this.chunkTokens) {
+      throw new Error(`chunkOverlapTokens (${this.chunkOverlapTokens}) must be smaller than ` + `chunkTokens (${this.chunkTokens})`);
+    }
+  }
+  async load() {
+    if (this._model)
+      return;
+    if (this._loadPromise)
+      return this._loadPromise;
+    this._loadPromise = (async () => {
+      const mod = await loadTransformersJs2();
+      let preferredDevice = this.device;
+      if (this.device === "auto") {
+        const nav = typeof navigator !== "undefined" ? navigator : null;
+        preferredDevice = "wasm";
+        if (nav?.gpu?.requestAdapter) {
+          try {
+            const adapter = await nav.gpu.requestAdapter();
+            if (adapter)
+              preferredDevice = "webgpu";
+          } catch {}
+        }
+      }
+      const tokenizer = await mod.AutoTokenizer.from_pretrained(this.model);
+      const model = await mod.AutoModelForTokenClassification.from_pretrained(this.model, {
+        device: preferredDevice,
+        dtype: this.dtype,
+        progress_callback: this._onProgress ? (raw) => this._emitProgress(raw, preferredDevice) : undefined
+      });
+      this._tokenizer = tokenizer;
+      this._model = model;
+      this._id2label = model.config.id2label;
+      const label2id = model.config.label2id;
+      this._oId = label2id?.["O"] ?? 0;
+      this.actualDevice = preferredDevice;
+    })();
+    return this._loadPromise;
+  }
+  async detect(text, opts) {
+    if (opts?.model && opts.model !== this.model) {
+      throw new Error(`CalibratedDetector was constructed with model='${this.model}'; ` + `per-call model override ('${opts.model}') not supported.`);
+    }
+    if (!text)
+      return [];
+    await this.load();
+    return this._detectChunked(text);
+  }
+  async _detectChunked(text) {
+    if (text.length <= 1500) {
+      return await this._detectWindow(text, 0);
+    }
+    const stride = 1500 - 200;
+    const seen = new Set;
+    const out = [];
+    for (let s = 0;s < text.length; s += stride) {
+      const e = Math.min(s + 1500, text.length);
+      const sub = text.slice(s, e);
+      const subSpans = await this._detectWindow(sub, s);
+      for (const sp of subSpans) {
+        const key = `${sp.start}|${sp.end}|${sp.label}`;
+        if (seen.has(key))
+          continue;
+        seen.add(key);
+        out.push(sp);
+      }
+      if (e >= text.length)
+        break;
+    }
+    out.sort((a, b) => a.start - b.start || a.end - b.end);
+    return out;
+  }
+  async _detectWindow(text, charOffset) {
+    const tokenizer = this._tokenizer;
+    const model = this._model;
+    const enc = await tokenizer([text], {
+      padding: true,
+      truncation: true,
+      max_length: this.chunkTokens + 2
+    });
+    const inputs = {};
+    if (enc.input_ids)
+      inputs["input_ids"] = enc.input_ids;
+    if (enc.attention_mask)
+      inputs["attention_mask"] = enc.attention_mask;
+    const out = await model(inputs);
+    const logits = readLogits(out.logits);
+    if (this.oBias !== 0) {
+      for (let t = 0;t < logits.length; t++) {
+        const row = logits[t];
+        if (row === undefined)
+          continue;
+        const cur = row[this._oId];
+        if (cur !== undefined) {
+          row[this._oId] = cur - this.oBias;
+        }
+      }
+    }
+    const inputIds = readInputIds(enc.input_ids);
+    const tokenLabels = [];
+    for (let t = 0;t < logits.length; t++) {
+      const row = logits[t];
+      const id = inputIds[t];
+      if (!row || id === undefined)
+        continue;
+      let maxIdx = 0;
+      let maxVal = row[0];
+      for (let i = 1;i < row.length; i++) {
+        if (row[i] > maxVal) {
+          maxVal = row[i];
+          maxIdx = i;
+        }
+      }
+      tokenLabels.push({ id, label: this._id2label[maxIdx] ?? "O" });
+    }
+    const spans = [];
+    let cursor = 0;
+    let groupIds = [];
+    let groupCat = null;
+    const flush = () => {
+      if (groupCat === null || groupIds.length === 0)
+        return;
+      const surface = normalizeDecoded(tokenizer.decode(groupIds, { skip_special_tokens: true }));
+      groupIds = [];
+      const cat = groupCat;
+      groupCat = null;
+      if (!surface)
+        return;
+      const idx = text.indexOf(surface, cursor);
+      if (idx === -1)
+        return;
+      let s = idx;
+      let e = idx + surface.length;
+      cursor = e;
+      if (this.trimWhitespace)
+        [s, e] = trimSpan2(text, s, e);
+      if (s >= e)
+        return;
+      spans.push({
+        start: s + charOffset,
+        end: e + charOffset,
+        label: cat,
+        text: text.slice(s, e),
+        score: 1
+      });
+    };
+    for (const { id, label } of tokenLabels) {
+      const cat = label === "O" ? null : label.split("-", 2)[1] ?? null;
+      if (cat === groupCat && cat !== null) {
+        groupIds.push(id);
+      } else {
+        flush();
+        if (cat !== null) {
+          groupIds.push(id);
+          groupCat = cat;
+        }
+      }
+    }
+    flush();
+    return spans;
+  }
+  _emitProgress(raw, device) {
+    if (!this._onProgress)
+      return;
+    const event = {
+      phase: raw.status ?? "progress",
+      device
+    };
+    if (raw.file)
+      event.file = raw.file;
+    if (raw.status === "progress") {
+      if (raw.loaded != null)
+        event.loadedBytes = raw.loaded;
+      if (raw.total != null)
+        event.totalBytes = raw.total;
+      if (raw.progress != null) {
+        const norm = raw.progress > 1 ? raw.progress / 100 : raw.progress;
+        event.fraction = Math.max(0, Math.min(1, norm));
+      }
+    }
+    this._onProgress(event);
+  }
+}
+async function loadTransformersJs2() {
+  try {
+    return await import("@huggingface/transformers");
+  } catch {
+    throw new Error("CalibratedDetector requires `@huggingface/transformers`. Install with: bun add @huggingface/transformers");
+  }
+}
+function readInputIds(tensor) {
+  if (!tensor)
+    return [];
+  if (typeof tensor.tolist === "function") {
+    const list = tensor.tolist();
+    if (Array.isArray(list) && Array.isArray(list[0])) {
+      return list[0].map((x) => Number(x));
+    }
+    if (Array.isArray(list))
+      return list.map((x) => Number(x));
+  }
+  const dims = tensor.dims;
+  const seq = dims[dims.length - 1] ?? 0;
+  const data = tensor.data;
+  const out = [];
+  for (let i = 0;i < seq; i++)
+    out.push(Number(data[i]));
+  return out;
+}
+function normalizeDecoded(s) {
+  return s.replace(/##/g, "").trim();
+}
+function readLogits(tensor) {
+  if (typeof tensor.tolist === "function") {
+    const list = tensor.tolist();
+    if (Array.isArray(list) && Array.isArray(list[0]) && Array.isArray(list[0][0])) {
+      return list[0];
+    }
+    if (Array.isArray(list) && Array.isArray(list[0])) {
+      return list;
+    }
+  }
+  const dims = tensor.dims;
+  const seq = dims[dims.length - 2] ?? 0;
+  const cls = dims[dims.length - 1] ?? 0;
+  const data = tensor.data;
+  const out = [];
+  for (let t = 0;t < seq; t++) {
+    const row = [];
+    for (let c = 0;c < cls; c++) {
+      row.push(Number(data[t * cls + c]));
+    }
+    out.push(row);
+  }
+  return out;
+}
 // src/detectors/index.ts
 function defaultDetector() {
   const hasKey = typeof process !== "undefined" && Boolean(process.env?.GHEIM_API_KEY);
-  return hasKey ? new RemoteDetector : new LocalDetector;
+  return hasKey ? new RemoteDetector : new CalibratedDetector;
 }
 // src/plain.ts
 function resolve(session, detector) {
@@ -896,4 +1190,4 @@ async function* deanonymizeStream(chunks, session) {
 }
 
 // src/index.ts
-var VERSION = "0.1.4";
+var VERSION = "0.1.6";

package/dist/cjs/openai/index.cjs

@@ -751,10 +751,281 @@ class CompositeDetector {
     return out;
   }
 }
+// src/detectors/calibrated.ts
+var DEFAULT_MODEL3 = "joelbarmettler/gheim-ch-560m";
+var DEFAULT_CHUNK_TOKENS2 = 510;
+var DEFAULT_OVERLAP_TOKENS2 = 32;
+function trimSpan2(text, start, end) {
+  while (start < end && /\s/.test(text[start] ?? ""))
+    start++;
+  while (end > start && /\s/.test(text[end - 1] ?? ""))
+    end--;
+  return [start, end];
+}
+
+class CalibratedDetector {
+  model;
+  oBias;
+  device;
+  dtype;
+  chunkTokens;
+  chunkOverlapTokens;
+  trimWhitespace;
+  actualDevice = null;
+  _model = null;
+  _tokenizer = null;
+  _id2label = {};
+  _oId = 0;
+  _loadPromise = null;
+  _onProgress;
+  constructor(opts = {}) {
+    this.model = opts.model ?? DEFAULT_MODEL3;
+    this.oBias = opts.oBias ?? 0.5;
+    this.device = opts.device ?? "auto";
+    this.dtype = opts.dtype ?? "fp32";
+    this.chunkTokens = opts.chunkTokens ?? DEFAULT_CHUNK_TOKENS2;
+    this.chunkOverlapTokens = opts.chunkOverlapTokens ?? DEFAULT_OVERLAP_TOKENS2;
+    this.trimWhitespace = opts.trimWhitespace ?? true;
+    this._onProgress = opts.onProgress;
+    if (this.chunkOverlapTokens >= this.chunkTokens) {
+      throw new Error(`chunkOverlapTokens (${this.chunkOverlapTokens}) must be smaller than ` + `chunkTokens (${this.chunkTokens})`);
+    }
+  }
+  async load() {
+    if (this._model)
+      return;
+    if (this._loadPromise)
+      return this._loadPromise;
+    this._loadPromise = (async () => {
+      const mod = await loadTransformersJs2();
+      let preferredDevice = this.device;
+      if (this.device === "auto") {
+        const nav = typeof navigator !== "undefined" ? navigator : null;
+        preferredDevice = "wasm";
+        if (nav?.gpu?.requestAdapter) {
+          try {
+            const adapter = await nav.gpu.requestAdapter();
+            if (adapter)
+              preferredDevice = "webgpu";
+          } catch {}
+        }
+      }
+      const tokenizer = await mod.AutoTokenizer.from_pretrained(this.model);
+      const model = await mod.AutoModelForTokenClassification.from_pretrained(this.model, {
+        device: preferredDevice,
+        dtype: this.dtype,
+        progress_callback: this._onProgress ? (raw) => this._emitProgress(raw, preferredDevice) : undefined
+      });
+      this._tokenizer = tokenizer;
+      this._model = model;
+      this._id2label = model.config.id2label;
+      const label2id = model.config.label2id;
+      this._oId = label2id?.["O"] ?? 0;
+      this.actualDevice = preferredDevice;
+    })();
+    return this._loadPromise;
+  }
+  async detect(text, opts) {
+    if (opts?.model && opts.model !== this.model) {
+      throw new Error(`CalibratedDetector was constructed with model='${this.model}'; ` + `per-call model override ('${opts.model}') not supported.`);
+    }
+    if (!text)
+      return [];
+    await this.load();
+    return this._detectChunked(text);
+  }
+  async _detectChunked(text) {
+    if (text.length <= 1500) {
+      return await this._detectWindow(text, 0);
+    }
+    const stride = 1500 - 200;
+    const seen = new Set;
+    const out = [];
+    for (let s = 0;s < text.length; s += stride) {
+      const e = Math.min(s + 1500, text.length);
+      const sub = text.slice(s, e);
+      const subSpans = await this._detectWindow(sub, s);
+      for (const sp of subSpans) {
+        const key = `${sp.start}|${sp.end}|${sp.label}`;
+        if (seen.has(key))
+          continue;
+        seen.add(key);
+        out.push(sp);
+      }
+      if (e >= text.length)
+        break;
+    }
+    out.sort((a, b) => a.start - b.start || a.end - b.end);
+    return out;
+  }
+  async _detectWindow(text, charOffset) {
+    const tokenizer = this._tokenizer;
+    const model = this._model;
+    const enc = await tokenizer([text], {
+      padding: true,
+      truncation: true,
+      max_length: this.chunkTokens + 2
+    });
+    const inputs = {};
+    if (enc.input_ids)
+      inputs["input_ids"] = enc.input_ids;
+    if (enc.attention_mask)
+      inputs["attention_mask"] = enc.attention_mask;
+    const out = await model(inputs);
+    const logits = readLogits(out.logits);
+    if (this.oBias !== 0) {
+      for (let t = 0;t < logits.length; t++) {
+        const row = logits[t];
+        if (row === undefined)
+          continue;
+        const cur = row[this._oId];
+        if (cur !== undefined) {
+          row[this._oId] = cur - this.oBias;
+        }
+      }
+    }
+    const inputIds = readInputIds(enc.input_ids);
+    const tokenLabels = [];
+    for (let t = 0;t < logits.length; t++) {
+      const row = logits[t];
+      const id = inputIds[t];
+      if (!row || id === undefined)
+        continue;
+      let maxIdx = 0;
+      let maxVal = row[0];
+      for (let i = 1;i < row.length; i++) {
+        if (row[i] > maxVal) {
+          maxVal = row[i];
+          maxIdx = i;
+        }
+      }
+      tokenLabels.push({ id, label: this._id2label[maxIdx] ?? "O" });
+    }
+    const spans = [];
+    let cursor = 0;
+    let groupIds = [];
+    let groupCat = null;
+    const flush = () => {
+      if (groupCat === null || groupIds.length === 0)
+        return;
+      const surface = normalizeDecoded(tokenizer.decode(groupIds, { skip_special_tokens: true }));
+      groupIds = [];
+      const cat = groupCat;
+      groupCat = null;
+      if (!surface)
+        return;
+      const idx = text.indexOf(surface, cursor);
+      if (idx === -1)
+        return;
+      let s = idx;
+      let e = idx + surface.length;
+      cursor = e;
+      if (this.trimWhitespace)
+        [s, e] = trimSpan2(text, s, e);
+      if (s >= e)
+        return;
+      spans.push({
+        start: s + charOffset,
+        end: e + charOffset,
+        label: cat,
+        text: text.slice(s, e),
+        score: 1
+      });
+    };
+    for (const { id, label } of tokenLabels) {
+      const cat = label === "O" ? null : label.split("-", 2)[1] ?? null;
+      if (cat === groupCat && cat !== null) {
+        groupIds.push(id);
+      } else {
+        flush();
+        if (cat !== null) {
+          groupIds.push(id);
+          groupCat = cat;
+        }
+      }
+    }
+    flush();
+    return spans;
+  }
+  _emitProgress(raw, device) {
+    if (!this._onProgress)
+      return;
+    const event = {
+      phase: raw.status ?? "progress",
+      device
+    };
+    if (raw.file)
+      event.file = raw.file;
+    if (raw.status === "progress") {
+      if (raw.loaded != null)
+        event.loadedBytes = raw.loaded;
+      if (raw.total != null)
+        event.totalBytes = raw.total;
+      if (raw.progress != null) {
+        const norm = raw.progress > 1 ? raw.progress / 100 : raw.progress;
+        event.fraction = Math.max(0, Math.min(1, norm));
+      }
+    }
+    this._onProgress(event);
+  }
+}
+async function loadTransformersJs2() {
+  try {
+    return await import("@huggingface/transformers");
+  } catch {
+    throw new Error("CalibratedDetector requires `@huggingface/transformers`. Install with: bun add @huggingface/transformers");
+  }
+}
+function readInputIds(tensor) {
+  if (!tensor)
+    return [];
+  if (typeof tensor.tolist === "function") {
+    const list = tensor.tolist();
+    if (Array.isArray(list) && Array.isArray(list[0])) {
+      return list[0].map((x) => Number(x));
+    }
+    if (Array.isArray(list))
+      return list.map((x) => Number(x));
+  }
+  const dims = tensor.dims;
+  const seq = dims[dims.length - 1] ?? 0;
+  const data = tensor.data;
+  const out = [];
+  for (let i = 0;i < seq; i++)
+    out.push(Number(data[i]));
+  return out;
+}
+function normalizeDecoded(s) {
+  return s.replace(/##/g, "").trim();
+}
+function readLogits(tensor) {
+  if (typeof tensor.tolist === "function") {
+    const list = tensor.tolist();
+    if (Array.isArray(list) && Array.isArray(list[0]) && Array.isArray(list[0][0])) {
+      return list[0];
+    }
+    if (Array.isArray(list) && Array.isArray(list[0])) {
+      return list;
+    }
+  }
+  const dims = tensor.dims;
+  const seq = dims[dims.length - 2] ?? 0;
+  const cls = dims[dims.length - 1] ?? 0;
+  const data = tensor.data;
+  const out = [];
+  for (let t = 0;t < seq; t++) {
+    const row = [];
+    for (let c = 0;c < cls; c++) {
+      row.push(Number(data[t * cls + c]));
+    }
+    out.push(row);
+  }
+  return out;
+}
 // src/detectors/index.ts
 function defaultDetector() {
   const hasKey = typeof process !== "undefined" && Boolean(process.env?.GHEIM_API_KEY);
-  return hasKey ? new RemoteDetector : new LocalDetector;
+  return hasKey ? new RemoteDetector : new CalibratedDetector;
 }
 
 // src/core/session.ts
@@ -815,8 +1086,8 @@ class Session {
     }
     out.push(text.slice(cursor));
     let redacted = out.join("");
-    redacted = recoverDuplicateOccurrences(redacted, this.mapping);
-    return { redacted, merged };
+    const recovered = recoverDuplicateOccurrences(text, redacted, merged, this.mapping);
+    return { redacted: recovered.redacted, merged: recovered.merged };
   }
   restore(text) {
     if (Object.keys(this.mapping).length === 0)
@@ -875,14 +1146,36 @@ function mergeAdjacent(text, spans) {
 function escapeRegex(s) {
   return s.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
 }
-function recoverDuplicateOccurrences(redacted, mapping) {
+function recoverDuplicateOccurrences(text, redacted, merged, mapping) {
+  const surfaceToLabel = new Map;
+  for (const m of merged)
+    surfaceToLabel.set(text.slice(m.start, m.end), m.label);
+  const covered = merged.map((m) => [m.start, m.end]);
+  const overlaps = (s, e) => covered.some(([cs, ce]) => s < ce && e > cs);
   const entries = Object.entries(mapping).filter(([, surface]) => surface && surface.length > 0).sort((a, b) => b[1].length - a[1].length);
   let result = redacted;
+  const additional = [];
   for (const [sentinel, surface] of entries) {
-    const pattern = new RegExp(`(?<![\\p{L}\\p{N}_])${escapeRegex(surface)}(?![\\p{L}\\p{N}_])`, "gu");
-    result = result.replace(pattern, sentinel);
+    const boundaryPattern = `(?<![\\p{L}\\p{N}_])${escapeRegex(surface)}(?![\\p{L}\\p{N}_])`;
+    result = result.replace(new RegExp(boundaryPattern, "gu"), sentinel);
+    const label = surfaceToLabel.get(surface);
+    if (!label)
+      continue;
+    const origRe = new RegExp(boundaryPattern, "gu");
+    let match;
+    while ((match = origRe.exec(text)) !== null) {
+      const start = match.index;
+      const end = start + surface.length;
+      if (overlaps(start, end))
+        continue;
+      additional.push({ start, end, label });
+      covered.push([start, end]);
+    }
   }
-  return result;
+  if (additional.length === 0)
+    return { redacted: result, merged };
+  const augmented = [...merged, ...additional].sort((a, b) => a.start - b.start);
+  return { redacted: result, merged: augmented };
 }
 
 // src/openai/_base.ts