Skip to content

pelinker.entity_head

Entity-assignment heads trained on HDBSCAN cluster labels (compact predict path).

EntityHead dataclass

Multi-class head mapping UMAP coords → emergent cluster id + confidence score.

Source code in pelinker/entity_head.py 
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
@dataclass
class EntityHead:
    """Multi-class head mapping UMAP coords → emergent cluster id + confidence score."""

    kind: str
    _estimator: MLPClassifier | LinearSVC
    classes_: np.ndarray

    def predict(self, X: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        """
        Predict cluster labels and per-row scores.

        Returns:
            ``(cluster_ids, scores)`` each shape ``(n_samples,)``. Scores are in ``[0, 1]``
            (``max(predict_proba)`` for MLP; logistic of max decision margin for LinearSVC).
        """
        if X.ndim != 2:
            raise ValueError(f"X must be 2D, got shape {X.shape}")
        xf = np.asarray(X, dtype=np.float64)
        if isinstance(self._estimator, MLPClassifier):
            proba = self._estimator.predict_proba(xf)
            idx = np.argmax(proba, axis=1)
            scores = proba[np.arange(len(idx)), idx].astype(np.float64, copy=False)
            labels = self._estimator.classes_[idx].astype(np.int64, copy=False)
            return labels, scores

        # LinearSVC: one-vs-rest decision margins → soft scores via logistic of max margin.
        pred = self._estimator.predict(xf)
        labels = np.asarray(pred, dtype=np.int64).ravel()
        decision = np.asarray(self._estimator.decision_function(xf), dtype=np.float64)
        if decision.ndim == 1:
            margin = decision
        else:
            margin = decision.max(axis=1)
        scores = 1.0 / (1.0 + np.exp(-margin))
        return labels, scores.astype(np.float64, copy=False)

predict(X)

Predict cluster labels and per-row scores.

Returns:

Type Description
ndarray

(cluster_ids, scores) each shape (n_samples,). Scores are in [0, 1]
ndarray

(max(predict_proba) for MLP; logistic of max decision margin for LinearSVC).
Source code in pelinker/entity_head.py 
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
def predict(self, X: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """
    Predict cluster labels and per-row scores.

    Returns:
        ``(cluster_ids, scores)`` each shape ``(n_samples,)``. Scores are in ``[0, 1]``
        (``max(predict_proba)`` for MLP; logistic of max decision margin for LinearSVC).
    """
    if X.ndim != 2:
        raise ValueError(f"X must be 2D, got shape {X.shape}")
    xf = np.asarray(X, dtype=np.float64)
    if isinstance(self._estimator, MLPClassifier):
        proba = self._estimator.predict_proba(xf)
        idx = np.argmax(proba, axis=1)
        scores = proba[np.arange(len(idx)), idx].astype(np.float64, copy=False)
        labels = self._estimator.classes_[idx].astype(np.int64, copy=False)
        return labels, scores

    # LinearSVC: one-vs-rest decision margins → soft scores via logistic of max margin.
    pred = self._estimator.predict(xf)
    labels = np.asarray(pred, dtype=np.int64).ravel()
    decision = np.asarray(self._estimator.decision_function(xf), dtype=np.float64)
    if decision.ndim == 1:
        margin = decision
    else:
        margin = decision.max(axis=1)
    scores = 1.0 / (1.0 + np.exp(-margin))
    return labels, scores.astype(np.float64, copy=False)

fit_linear_svc_entity_head(X_umap, cluster_labels, *, random_state=13)

Fit a linear multi-class SVM (study underfitting control; not the shipped default).

Source code in pelinker/entity_head.py 
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
def fit_linear_svc_entity_head(
    X_umap: np.ndarray,
    cluster_labels: np.ndarray,
    *,
    random_state: int = 13,
) -> EntityHead:
    """Fit a linear multi-class SVM (study underfitting control; not the shipped default)."""
    X, y = _non_noise_xy(X_umap, cluster_labels)
    est = LinearSVC(max_iter=5000, random_state=random_state)
    est.fit(X, y)
    return EntityHead(
        kind="linear_svc",
        _estimator=est,
        classes_=np.asarray(est.classes_, dtype=np.int64),
    )

fit_mlp_entity_head(X_umap, cluster_labels, *, hidden_layer_sizes=DEFAULT_ENTITY_HEAD_HIDDEN_LAYERS, random_state=13)

Fit an MLP on non-noise HDBSCAN labels (cluster_labels != -1).

Source code in pelinker/entity_head.py 
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
def fit_mlp_entity_head(
    X_umap: np.ndarray,
    cluster_labels: np.ndarray,
    *,
    hidden_layer_sizes: tuple[int, ...] = DEFAULT_ENTITY_HEAD_HIDDEN_LAYERS,
    random_state: int = 13,
) -> EntityHead:
    """Fit an MLP on non-noise HDBSCAN labels (``cluster_labels != -1``)."""
    X, y = _non_noise_xy(X_umap, cluster_labels)
    est = MLPClassifier(
        hidden_layer_sizes=hidden_layer_sizes,
        activation="relu",
        solver="adam",
        early_stopping=True,
        validation_fraction=0.1,
        max_iter=500,
        random_state=random_state,
    )
    est.fit(X, y)
    return EntityHead(
        kind="mlp",
        _estimator=est,
        classes_=np.asarray(est.classes_, dtype=np.int64),
    )