| |
| |
| """Utterance-convention decoder: pure-numpy predictor for exported gradient-boosted trees, plus the decision rule. |
| |
| The trees are scikit-learn HistGradientBoostingRegressor models exported by tools/export_decoder.py. Prediction |
| replicates sklearn exactly: float64 baseline plus each tree's leaf value, added in iteration order; numerical splits |
| go left when x <= threshold, missing values follow the stored missing direction. |
| Decision rule (as frozen in the research pipeline): bridge a same-speaker gap when the gap model predicts a negative |
| error change per second; drop a short segment when the segment model predicts a negative error change. |
| """ |
| import json |
|
|
| import numpy as np |
|
|
| from .geometry import union |
|
|
|
|
| class TreeEnsemble: |
| def __init__(self, blob): |
| self.features = blob["features"] |
| self.baseline = float(blob["baseline"]) |
| self.trees = [{k: np.asarray(v) for k, v in t.items()} for t in blob["trees"]] |
|
|
| def predict(self, rows): |
| if not rows: |
| return np.zeros(0) |
| X = np.array([[r[f] for f in self.features] for r in rows], dtype=np.float64) |
| out = np.full(len(X), self.baseline, dtype=np.float64) |
| for t in self.trees: |
| feat, thr, mleft = t["feature"], t["threshold"], t["missing_left"] |
| left, right, leaf, value = t["left"], t["right"], t["leaf"], t["value"] |
| node = np.zeros(len(X), dtype=np.int64) |
| for _ in range(len(feat)): |
| active = ~leaf[node] |
| if not active.any(): |
| break |
| n = node[active] |
| x = X[np.flatnonzero(active), feat[n]] |
| go_left = np.where(np.isnan(x), mleft[n], x <= thr[n]) |
| node[active] = np.where(go_left, left[n], right[n]) |
| out += value[node] |
| return out |
|
|
|
|
| def load_decoder(path): |
| blob = json.loads(open(path).read()) |
| return TreeEnsemble(blob["gap"]), TreeEnsemble(blob["seg"]) |
|
|
|
|
| def apply(hyp, gap_rows, gap_pred, seg_rows, seg_pred, tau_gap=0.0, tau_seg=0.0): |
| """hyp: spk -> unbridged intervals. Bridge gaps with gap_pred > tau_gap; drop segments with seg_pred < tau_seg |
| unless they were merged by a bridge. (Research: convention_decoder.apply without the unused `moves` option.)""" |
| fill = {(r["spk"], r["j"]) for r, p in zip(gap_rows, gap_pred) if p > tau_gap} |
| drop = {(r["spk"], r["j"]) for r, p in zip(seg_rows, seg_pred) if p < tau_seg} |
| out = {} |
| for k, v in hyp.items(): |
| iv = union([tuple(x) for x in v]) |
| keep = [] |
| for j, (s, e) in enumerate(iv): |
| merged = (k, j) in fill or (k, j + 1) in fill |
| keep.append(not ((k, j) in drop and not merged)) |
| res = [] |
| for j, (s, e) in enumerate(iv): |
| if j > 0 and (k, j) in fill and res and keep[j] and keep[j - 1]: |
| res[-1][1] = max(res[-1][1], e) |
| elif keep[j]: |
| res.append([s, e]) |
| if res: |
| out[k] = [(a, b) for a, b in res] |
| return out |
|
|
