langset 0.1.0__py3-none-any.whl

langset/__init__.py

@@ -0,0 +1,10 @@
+"""langset — few-shot fine-tune an LLM to emit a latent into a bespoke, specialized geometry.
+
+The trained model is Sentence-Transformer-shaped, so it drops into SetFit as a `model_body`.
+"""
+from langset.modeling import EmitHead, LangSetModel
+from langset.trainer import Trainer
+from langset.training_args import TrainingArguments
+
+__all__ = ["LangSetModel", "EmitHead", "Trainer", "TrainingArguments"]
+__version__ = "0.1.0"

langset/data.py

@@ -0,0 +1,29 @@
+"""Dataset contract for langset.
+
+A training example is a row with:
+  - `input_text`   : the text you'll have at inference (e.g. a name, a query, a review)
+  - `target_text`  : a description of the SAME item that defines where it should land (embedded by the
+                     bootstrap encoder to seed the geometry). Point this at something `input_text` can't
+                     trivially regenerate, or you're just distilling a text encoder.
+  - any number of optional label columns : EVAL-ONLY geometry probes (kNN-purity at validation). They never
+                     enter training — training on them would collapse the geometry onto the labels.
+
+Pass a `datasets.Dataset` or a `list[dict]` to `Trainer`; use `column_mapping` to rename your columns onto
+`input_text` / `target_text`.
+"""
+from __future__ import annotations
+
+from typing import Any, Optional
+
+
+def from_records(records: list[dict[str, Any]], input_key: str, target_key: str,
+                 label_keys: Optional[list[str]] = None) -> list[dict[str, Any]]:
+    """Convenience: project a list of dicts onto the langset contract."""
+    label_keys = label_keys or []
+    out = []
+    for r in records:
+        row = {"input_text": str(r[input_key]), "target_text": str(r[target_key])}
+        for k in label_keys:
+            row[k] = r.get(k, "unknown")
+        out.append(row)
+    return out

langset/modeling.py

@@ -0,0 +1,163 @@
+"""LangSetModel: an LLM (LoRA) + a learned emit head that maps input text -> a latent in a bespoke geometry.
+
+The output is Sentence-Transformer-shaped (`encode`, `get_sentence_embedding_dimension`, `as_sentence_transformer`)
+so the trained model drops straight into SetFit as a `model_body`.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, Optional, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class EmitHead(nn.Module):
+    """Learned query tokens whose post-backbone hidden states are read out as the latent (the LLM's vector 'mouth').
+
+    `dropout` is trained ON when >0 so MC-dropout could later serve as an uncertainty signal; it's also plain
+    regularization. The head stays fp32 even when the backbone is bf16.
+    """
+
+    def __init__(self, h: int, d: int, n_emit: int = 4, dropout: float = 0.0) -> None:
+        super().__init__()
+        self.n_emit = n_emit
+        self.q = nn.Parameter(torch.randn(n_emit, h) * 0.02)
+        self.drop = nn.Dropout(dropout)
+        self.out = nn.Linear(n_emit * h, d)
+
+    def forward(self, hid_emit: torch.Tensor) -> torch.Tensor:           # [B, n_emit, h] -> [B, d]
+        flat = self.drop(hid_emit.reshape(hid_emit.size(0), -1).float())
+        return F.normalize(self.out(flat), p=2, dim=-1)
+
+
+def build_backbone(llm_model: str, lora_r: int, dropout: float, bf16: bool, dev: str) -> Any:
+    from peft import LoraConfig, get_peft_model  # type: ignore[import-untyped]
+    from transformers import AutoModelForCausalLM  # type: ignore[import-untyped]
+    dt = torch.bfloat16 if bf16 else torch.float32
+    base = AutoModelForCausalLM.from_pretrained(
+        llm_model, torch_dtype=dt, attention_dropout=dropout, attn_implementation="sdpa").to(dev)
+    lora = LoraConfig(r=lora_r, lora_alpha=2 * lora_r, lora_dropout=dropout, bias="none",
+                      target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
+                                      "gate_proj", "up_proj", "down_proj"])
+    return get_peft_model(base, lora)
+
+
+class LangSetModel(nn.Module):
+    """LLM backbone (LoRA) + EmitHead. `from_pretrained` infers the latent dim from the bootstrap encoder, so the
+    emitted geometry starts in that encoder's space (the Trainer then EMA-specializes it away)."""
+
+    def __init__(self, backbone: Any, tokenizer: Any, latent_dim: int, n_emit: int,
+                 llm_model: str, bootstrap_model: str, dropout: float = 0.0, max_len: int = 512) -> None:
+        super().__init__()
+        self.backbone = backbone
+        self.tokenizer = tokenizer
+        self.embed = backbone.get_input_embeddings()
+        self.h = int(backbone.config.hidden_size)
+        self.latent_dim = latent_dim
+        self.head = EmitHead(self.h, latent_dim, n_emit, dropout)
+        self.llm_model = llm_model
+        self.bootstrap_model = bootstrap_model
+        self.max_len = max_len
+        self._bootstrap: Any = None
+
+    # ---- construction ----
+    @classmethod
+    def from_pretrained(cls, llm_model: str, bootstrap_model: str, *, latent_dim: Optional[int] = None,
+                        n_emit: int = 4, lora_r: int = 16, dropout: float = 0.0, bf16: bool = False,
+                        max_len: int = 512, device: Optional[str] = None) -> "LangSetModel":
+        from sentence_transformers import SentenceTransformer  # type: ignore[import-untyped]
+        from transformers import AutoTokenizer  # type: ignore[import-untyped]
+        dev = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        tok = AutoTokenizer.from_pretrained(llm_model)
+        if tok.pad_token_id is None:
+            tok.pad_token = tok.eos_token
+        if latent_dim is None:                                   # target space dim = bootstrap encoder dim
+            _st = SentenceTransformer(bootstrap_model)            # ST renamed the method in 4.x -> support both
+            _dim = getattr(_st, "get_embedding_dimension", None) or _st.get_sentence_embedding_dimension
+            latent_dim = int(_dim())
+        backbone = build_backbone(llm_model, lora_r, dropout, bf16, dev)
+        m = cls(backbone, tok, latent_dim, n_emit, llm_model, bootstrap_model, dropout, max_len).to(dev)
+        return m
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.backbone.parameters()).device
+
+    @property
+    def bootstrap_encoder(self) -> Any:
+        """Lazily-loaded target encoder (used by the Trainer for targets + the beats-bootstrap baseline)."""
+        if self._bootstrap is None:
+            from sentence_transformers import SentenceTransformer  # type: ignore[import-untyped]
+            self._bootstrap = SentenceTransformer(self.bootstrap_model, device=str(self.device))
+        return self._bootstrap
+
+    # ---- forward / inference ----
+    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
+        rev = self.embed(input_ids)
+        q = self.head.q.unsqueeze(0).expand(input_ids.size(0), -1, -1).to(rev.dtype)
+        emb = torch.cat([rev, q], 1)
+        am = torch.cat([attention_mask,
+                        torch.ones(input_ids.size(0), self.head.n_emit, device=input_ids.device,
+                                   dtype=attention_mask.dtype)], 1)
+        hid = self.backbone(inputs_embeds=emb, attention_mask=am, output_hidden_states=True).hidden_states[-1]
+        return self.head(hid[:, -self.head.n_emit:, :])
+
+    @torch.no_grad()
+    def encode(self, sentences: Union[str, list[str]], batch_size: int = 32, convert_to_numpy: bool = True,
+               normalize_embeddings: bool = True, show_progress_bar: bool = False,
+               device: Optional[str] = None) -> Union[np.ndarray, torch.Tensor]:
+        """Sentence-Transformer-compatible. This is the method SetFit calls on its body."""
+        single = isinstance(sentences, str)
+        texts = [sentences] if single else list(sentences)
+        was_training = self.training
+        self.eval()
+        out: list[torch.Tensor] = []
+        for i in range(0, len(texts), batch_size):
+            enc = self.tokenizer(texts[i:i + batch_size], padding=True, truncation=True,
+                                 max_length=self.max_len, return_tensors="pt").to(self.device)
+            z = self(enc["input_ids"], enc["attention_mask"])
+            if normalize_embeddings:
+                z = F.normalize(z, p=2, dim=-1)
+            out.append(z.cpu())
+        if was_training:
+            self.train()
+        emb = torch.cat(out)
+        emb = emb[0] if single else emb
+        return emb.numpy() if convert_to_numpy else emb
+
+    def emit(self, sentences: Union[str, list[str]], **kw: Any) -> torch.Tensor:
+        return self.encode(sentences, convert_to_numpy=False, **kw)  # type: ignore[return-value]
+
+    def get_sentence_embedding_dimension(self) -> int:
+        return self.latent_dim
+
+    def as_sentence_transformer(self) -> Any:
+        """Wrap as a `sentence_transformers.SentenceTransformer` so it drops into SetFit as `model_body`."""
+        from langset.st_module import to_sentence_transformer
+        return to_sentence_transformer(self)
+
+    # ---- persistence (LoRA + head + config; backbone/bootstrap rebuilt from ids) ----
+    def save_pretrained(self, path: Union[str, Path]) -> None:
+        import json
+        p = Path(path); p.mkdir(parents=True, exist_ok=True)
+        torch.save({"head": self.head.state_dict(),
+                    "lora": {k: v.cpu() for k, v in self.backbone.state_dict().items() if "lora" in k}},
+                   p / "langset.pt")
+        (p / "config.json").write_text(json.dumps({
+            "llm_model": self.llm_model, "bootstrap_model": self.bootstrap_model,
+            "latent_dim": self.latent_dim, "n_emit": self.head.n_emit, "max_len": self.max_len}))
+
+    @classmethod
+    def load(cls, path: Union[str, Path], *, lora_r: int = 16, device: Optional[str] = None) -> "LangSetModel":
+        import json
+        p = Path(path); cfg = json.loads((p / "config.json").read_text())
+        m = cls.from_pretrained(cfg["llm_model"], cfg["bootstrap_model"], latent_dim=cfg["latent_dim"],
+                                n_emit=cfg["n_emit"], lora_r=lora_r, max_len=cfg["max_len"], device=device)
+        sd = torch.load(p / "langset.pt", map_location=m.device, weights_only=False)
+        m.backbone.load_state_dict(sd["lora"], strict=False)
+        m.head.load_state_dict(sd["head"])
+        m.eval()
+        return m

langset/selection.py

@@ -0,0 +1,81 @@
+"""Validation / early-stop metrics. The hard-won rules:
+  - NEVER select on training loss (InfoNCE+EMA can minimize it by COLLAPSING the geometry).
+  - NEVER score retrieval against the frozen bootstrap targets (the model specializes AWAY from them).
+  - Score on held-out, in the CURRENT geometry: retrieval mrr + a collapse guard; and, when the rows carry
+    geometry labels, held-out kNN purity (+ beats-bootstrap), which is the preferred, intent-aligned selector.
+Geometry labels are EVAL-ONLY probes — they never enter training (training on them would collapse the space).
+"""
+from __future__ import annotations
+
+from typing import Any, Optional
+
+import numpy as np
+
+
+def retrieval_mrr(pred: np.ndarray, bank: np.ndarray) -> dict[str, float]:
+    """Each pred[i] should retrieve target bank[i] among the val set (rank-based, scale-free, collapse-sensitive)."""
+    sims = pred @ bank.T
+    truth = np.arange(len(pred))
+    order = np.argsort(-sims, axis=1)
+    rank = np.array([int(np.where(order[i] == truth[i])[0][0]) for i in range(len(truth))])
+    return {"r@1": float((rank == 0).mean()), "mrr": float((1.0 / (rank + 1)).mean())}
+
+
+def collapse_score(X: np.ndarray) -> float:
+    """Mean off-diagonal cosine of (normalized) emissions. ->1 means collapsed; healthy spaces sit low."""
+    Xn = X / (np.linalg.norm(X, axis=1, keepdims=True) + 1e-9)
+    s = Xn @ Xn.T
+    np.fill_diagonal(s, np.nan)
+    return float(np.nanmean(s))
+
+
+def knn_purity(X: np.ndarray, labels: list[str], k: int = 5) -> float:
+    """Leave-one-out majority-vote accuracy by label (does the geometry organize by this attribute?)."""
+    keep = [i for i, l in enumerate(labels) if str(l) not in ("", "unknown", "none", "nan")]
+    if len(keep) < k + 2:
+        return float("nan")
+    Xk = X[keep] / (np.linalg.norm(X[keep], axis=1, keepdims=True) + 1e-9)
+    lk = [labels[i] for i in keep]
+    sims = Xk @ Xk.T
+    np.fill_diagonal(sims, -1e9)
+    nn = np.argsort(-sims, axis=1)[:, :k]
+    correct = sum(max(set(v := [lk[j] for j in row]), key=v.count) == lk[i] for i, row in enumerate(nn))
+    return correct / len(keep)
+
+
+def evaluate(emit_val: np.ndarray, target_val: np.ndarray, label_cols: dict[str, list[str]],
+             bootstrap_val: Optional[np.ndarray], select: str) -> dict[str, Any]:
+    """Compute all val signals + the scalar `score` used for selection/early-stop."""
+    out: dict[str, Any] = {}
+    out.update(retrieval_mrr(emit_val, target_val))
+    out["collapse"] = collapse_score(emit_val)
+    have_labels = len(label_cols) > 0
+    if have_labels:
+        beaten, purities = 0, []
+        for name, labels in label_cols.items():
+            pe = knn_purity(emit_val, labels)
+            out[f"purity/{name}"] = pe
+            if pe == pe:
+                purities.append(pe)
+                if bootstrap_val is not None:
+                    pb = knn_purity(bootstrap_val, labels)
+                    out[f"purity_bootstrap/{name}"] = pb
+                    beaten += int(pe > pb)
+        out["purity_mean"] = float(np.mean(purities)) if purities else float("nan")
+        out["beats_bootstrap"] = beaten
+
+    mode = select
+    if mode == "auto":
+        mode = "purity" if have_labels else "retrieval"
+    if mode == "purity" and have_labels:
+        score = out.get("beats_bootstrap", 0) + (out.get("purity_mean", 0.0) or 0.0)
+    elif mode == "loss":
+        score = -out["collapse"]                       # discouraged; at least penalize collapse
+    else:
+        score = out["mrr"]
+    # collapse guard: a fully collapsed geometry is never "best"
+    if out["collapse"] > 0.95:
+        score = -1.0
+    out["score"] = float(score)
+    out["select_mode"] = mode
+    return out

langset/st_module.py

@@ -0,0 +1,62 @@
+"""Wrap a trained LangSetModel as a `sentence_transformers.SentenceTransformer` so it drops into SetFit as the
+`model_body`:
+
+    body = langset_model.as_sentence_transformer()
+    setfit_model = SetFitModel(model_body=body, labels=[...])
+"""
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+from torch import nn
+
+
+class LangSetSTModule(nn.Module):
+    """A Sentence-Transformers custom module: tokenizes text and emits the langset latent as `sentence_embedding`."""
+
+    def __init__(self, model: Any) -> None:
+        super().__init__()
+        self.model = model                                  # registered submodule -> .to(device) propagates
+        # attributes SentenceTransformer / SetFit's trainer expect a body to expose
+        self.tokenizer = model.tokenizer
+        self.auto_model = model.backbone
+        self.max_seq_length = model.max_len
+
+    def forward(self, features: dict[str, Any]) -> dict[str, Any]:
+        z = self.model(features["input_ids"], features["attention_mask"])
+        features["sentence_embedding"] = z
+        return features
+
+    def tokenize(self, texts: list[str], **kw: Any) -> dict[str, torch.Tensor]:
+        enc = self.model.tokenizer(list(texts), padding=True, truncation=True,
+                                   max_length=self.model.max_len, return_tensors="pt")
+        return {"input_ids": enc["input_ids"], "attention_mask": enc["attention_mask"]}
+
+    def get_sentence_embedding_dimension(self) -> int:
+        return int(self.model.latent_dim)
+
+    def get_embedding_dimension(self) -> int:        # ST 4.x name
+        return int(self.model.latent_dim)
+
+    def get_config_dict(self) -> dict[str, Any]:
+        return {}
+
+    def save(self, output_path: str, **kw: Any) -> None:
+        self.model.save_pretrained(output_path)
+
+    @staticmethod
+    def load(input_path: str, **kw: Any) -> "LangSetSTModule":
+        from langset.modeling import LangSetModel
+        return LangSetSTModule(LangSetModel.load(input_path))
+
+
+def to_sentence_transformer(model: Any) -> Any:
+    from sentence_transformers import SentenceTransformer  # type: ignore[import-untyped]
+    st = SentenceTransformer(modules=[LangSetSTModule(model)], device=str(model.device))
+    for attr, val in (("tokenizer", model.tokenizer), ("max_seq_length", model.max_len)):
+        try:                                                 # some ST versions expect these on the ST object
+            setattr(st, attr, val)
+        except Exception:                                    # noqa: BLE001 - it's a read-only property -> already resolves
+            pass
+    return st

langset/trainer.py

@@ -0,0 +1,135 @@
+"""Trainer: bootstrap the target geometry from the bootstrap encoder, contrastively fit the LLM emitter, and
+EMA-specialize the geometry away from the seed. Selects/early-stops on held-out geometry (see selection.py),
+never on training loss.
+
+Dataset rows: `input_text`, `target_text`, plus optional geometry-label columns (EVAL-ONLY probes).
+"""
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, Optional, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from langset import selection
+from langset.modeling import LangSetModel
+from langset.training_args import TrainingArguments
+
+
+def _columns(dataset: Any) -> dict[str, list[Any]]:
+    if hasattr(dataset, "column_names"):                       # datasets.Dataset
+        return {c: list(dataset[c]) for c in dataset.column_names}
+    rows = list(dataset)                                       # list[dict]
+    return {k: [r[k] for r in rows] for k in rows[0]}
+
+
+class Trainer:
+    def __init__(self, model: LangSetModel, args: TrainingArguments, train_dataset: Any,
+                 eval_dataset: Optional[Any] = None, column_mapping: Optional[dict[str, str]] = None,
+                 label_columns: Optional[list[str]] = None) -> None:
+        self.model = model
+        self.args = args
+        cols = _columns(train_dataset)
+        cmap = column_mapping or {}
+        inv = {v: k for k, v in cmap.items()}                 # user-col -> canonical
+        get = lambda canon: cols[inv.get(canon, canon)]       # noqa: E731
+        self.input_text = [str(x) for x in get("input_text")]
+        self.target_text = [str(x) for x in get("target_text")]
+        used = {inv.get("input_text", "input_text"), inv.get("target_text", "target_text")}
+        lc = label_columns if label_columns is not None else [c for c in cols if c not in used]
+        self.label_cols = {c: [str(x) for x in cols[c]] for c in lc}
+        if args.verbose:
+            print(f"[langset] {len(self.input_text)} rows | label columns (eval-only): {list(self.label_cols)}",
+                  flush=True)
+
+    def train(self) -> LangSetModel:
+        a, m = self.args, self.model
+        dev = m.device
+        torch.manual_seed(a.seed)
+        rng = np.random.default_rng(a.seed)
+
+        # 1. bootstrap targets = bootstrap_encoder(target_text)  (the seed geometry)
+        tgt_np = m.bootstrap_encoder.encode(self.target_text, normalize_embeddings=True,
+                                            convert_to_numpy=True, show_progress_bar=False).astype(np.float32)
+        tgt = torch.tensor(tgt_np, device=dev)
+        n = len(self.input_text)
+        perm = rng.permutation(n)
+        n_val = max(4, int(n * a.val_frac))
+        val_idx = perm[:n_val]
+        tr_idx = perm[n_val:]
+
+        tok = m.tokenizer
+        enc_in = tok(self.input_text, padding=True, truncation=True, max_length=a.max_len, return_tensors="pt")
+        ids = enc_in["input_ids"].to(dev)
+        mask = enc_in["attention_mask"].to(dev)
+
+        ema = tgt.clone()                                     # EMA self-target (the "specialize" step)
+        opt = torch.optim.AdamW([p for p in m.parameters() if p.requires_grad], lr=a.lr)
+        run = None
+        if a.report_to == "wandb":
+            import wandb  # type: ignore[import-untyped]
+            run = wandb.init(project=a.wandb_project, config=vars(a))
+
+        best_score, best_state, no_improve = -1e9, None, 0
+        for ep in range(a.epochs):
+            m.train()
+            order = tr_idx[rng.permutation(len(tr_idx))]
+            tot = nb = 0.0
+            for i in range(0, len(order), a.batch_size):
+                idx = torch.tensor(order[i:i + a.batch_size], device=dev)
+                pred = m(ids[idx], mask[idx])
+                target = ema[idx] if a.ema else tgt[idx]
+                logits = (pred @ target.t()) / a.tau          # InfoNCE, in-batch negatives
+                loss = F.cross_entropy(logits, torch.arange(len(idx), device=dev))
+                if a.lam_anchor > 0:                          # pull toward the frozen bootstrap (0 = fully specialize)
+                    loss = loss + a.lam_anchor * (1 - (pred * tgt[idx]).sum(1)).mean()
+                opt.zero_grad(); loss.backward(); opt.step()
+                if a.ema:
+                    with torch.no_grad():
+                        upd = a.ema_m * ema[idx] + (1 - a.ema_m) * pred.detach()
+                        ema[idx] = F.normalize(upd, p=2, dim=-1)
+                tot += float(loss.detach()); nb += 1
+
+            if ep % a.eval_every:
+                continue
+            # 2. validate in the CURRENT geometry: input-view vs target-view co-location (never vs the seed)
+            vi = self.input_text and [self.input_text[j] for j in val_idx]
+            vt = [self.target_text[j] for j in val_idx]
+            emit_in = m.encode(vi, normalize_embeddings=True)
+            emit_tg = m.encode(vt, normalize_embeddings=True)
+            labels_val = {k: [v[j] for j in val_idx] for k, v in self.label_cols.items()}
+            metrics = selection.evaluate(emit_in, emit_tg, labels_val, tgt_np[val_idx], a.select)
+            if a.verbose:
+                extra = f" purity_mean={metrics.get('purity_mean'):.3f} beats={metrics.get('beats_bootstrap')}" \
+                    if "purity_mean" in metrics else ""
+                print(f"ep{ep:02d} loss={tot/nb:.3f} mrr={metrics['mrr']:.3f} collapse={metrics['collapse']:.3f}"
+                      f"{extra} score={metrics['score']:.3f} [{metrics['select_mode']}]", flush=True)
+            if run is not None:
+                run.log({"loss": tot / nb, **metrics, "epoch": ep})
+
+            if metrics["score"] > best_score:
+                best_score = metrics["score"]
+                best_state = {"head": {k: v.detach().cpu().clone() for k, v in m.head.state_dict().items()},
+                              "lora": {k: v.detach().cpu().clone()
+                                       for k, v in m.backbone.state_dict().items() if "lora" in k}}
+                no_improve = 0
+            else:
+                no_improve += 1
+                if no_improve >= a.patience:
+                    if a.verbose:
+                        print(f"[langset] early stop at ep{ep} (best score {best_score:.3f})", flush=True)
+                    break
+
+        if best_state is not None:                            # restore best
+            m.head.load_state_dict(best_state["head"])
+            m.backbone.load_state_dict(best_state["lora"], strict=False)
+        m.eval()
+        Path(a.output_dir).mkdir(parents=True, exist_ok=True)
+        m.save_pretrained(a.output_dir)
+        if run is not None:
+            run.finish()
+        if a.verbose:
+            print(f"[langset] done. best score={best_score:.3f} -> {a.output_dir}", flush=True)
+        return m

langset/training_args.py

@@ -0,0 +1,34 @@
+"""Training configuration for langset. Defaults encode the lessons from the validated recipe:
+EMA-specialize (ema_m + low anchor), select on geometry not loss, restore-best."""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Literal, Optional
+
+
+@dataclass
+class TrainingArguments:
+    # optimization
+    epochs: int = 40
+    batch_size: int = 16
+    lr: float = 3e-4
+    tau: float = 0.07                     # InfoNCE temperature
+    max_len: int = 512
+
+    # bootstrap -> specialize
+    ema: bool = True                      # EMA self-distillation target (the "specialize" step)
+    ema_m: float = 0.99                   # EMA momentum (lower = faster drift off the bootstrap)
+    lam_anchor: float = 0.1               # cosine pull toward the FROZEN bootstrap target (0 = fully specialize)
+
+    # validation / early-stop  (see selection.py)
+    val_frac: float = 0.2
+    eval_every: int = 1
+    patience: int = 6                     # epochs without val improvement before stopping
+    select: Literal["auto", "retrieval", "purity", "loss"] = "auto"  # auto = purity if labels else retrieval
+    seed: int = 0
+
+    # io / logging
+    output_dir: str = "langset-out"
+    report_to: Optional[str] = None       # "wandb" or None
+    wandb_project: str = "langset"
+    verbose: bool = True

langset-0.1.0.dist-info/METADATA

@@ -0,0 +1,172 @@
+Metadata-Version: 2.4
+Name: langset
+Version: 0.1.0
+Summary: Few-shot fine-tune an LLM to emit a latent into a bespoke, specialized geometry — and drop the result into SetFit as a Sentence-Transformer body.
+Author: Stephen Solka
+License: Apache-2.0
+License-File: LICENSE
+Requires-Python: >=3.10
+Requires-Dist: datasets>=2.15.0
+Requires-Dist: numpy>=1.26
+Requires-Dist: peft>=0.13.2
+Requires-Dist: sentence-transformers>=3.0
+Requires-Dist: torch>=2.2
+Requires-Dist: transformers>=4.41
+Provides-Extra: dev
+Requires-Dist: pytest; extra == 'dev'
+Requires-Dist: ruff; extra == 'dev'
+Provides-Extra: setfit
+Requires-Dist: scikit-learn; extra == 'setfit'
+Requires-Dist: setfit>=1.1.0; extra == 'setfit'
+Requires-Dist: torch<2.5,>=2.2; extra == 'setfit'
+Requires-Dist: transformers<4.47,>=4.41; extra == 'setfit'
+Provides-Extra: wandb
+Requires-Dist: wandb>=0.18; extra == 'wandb'
+Description-Content-Type: text/markdown
+
+# langset — read text, answer in a vector
+
+**langset turns a language model into your own bespoke embedding space, few-shot.** Bolt a tiny vector head
+onto a pretrained LLM, point it at a target geometry you define, and it learns to *read text and emit a latent*
+into that space — using the LLM's world knowledge to do the work. From ~480 examples it builds a specialized
+"how it sounds" geometry that **beats the encoder it was bootstrapped from on 4/5 held-out axes** — while the
+*same architecture with a randomly-initialized backbone sits at chance*. That gap is the world knowledge, and
+it's the whole point.
+
+What makes langset different:
+
+* 🎯 **Latent out, not a label.** You define the output space; the model answers in a *vector*, not a class.
+  The geometry can be anything a description can seed — "how it sounds", "how it behaves", "what it's like".
+* 🧠 **World knowledge does the work.** It's a generative LLM, so it generalizes from *hundreds* of examples,
+  not millions. Swap in a random-init backbone and it barely moves off chance — the pretrained knowledge is
+  the engine.
+* 🌱 **Bootstrap, then specialize.** Seed the target geometry from *any* off-the-shelf encoder (no contrastive
+  pairs, no labels), then an EMA self-target drifts it *off* that seed into your own task-shaped space — it
+  stops being "text similarity" and becomes its own thing (measurably: it leaves the seed's cone).
+* 🔀 **Input-agnostic & non-circular.** Many input views can map to the same point, and you point the target at
+  a signal the input text can't regenerate — so it's not just distilling a text encoder.
+* ✅ **Honest selection.** Optional per-row geometry labels are *eval-only* probes; langset early-stops on
+  their held-out kNN-purity with a collapse guard — never on the training loss (which collapse can game).
+
+## Install
+
+```bash
+pip install langset
+```
+
+## Usage
+
+A langset dataset is rows of `input_text` → `target_text` (+ optional eval-only geometry labels). You pick the
+LLM backbone and the bootstrap encoder; langset trains the mapping and specializes the geometry.
+
+```python
+from langset import LangSetModel, Trainer, TrainingArguments
+
+rows = [  # a review SNIPPET (what you'll have at inference) -> a description that defines where it should land
+    {"input_text": "an hour-long track of detuned riffs that never break stride, moving at the pace of continental drift",
+     "target_text": "glacial detuned doom-metal, sludgy and hypnotic", "mood": "heavy"},
+    {"input_text": "chopped vocal ghosts drifting over vinyl crackle and the hiss of a city at 3am",
+     "target_text": "crackly nocturnal UK garage, pitched vocal ghosts", "mood": "calm"},
+    # ...
+]
+
+model = LangSetModel.from_pretrained(
+    llm_model="HuggingFaceTB/SmolLM2-135M",                    # any HF causal LM (this is what the examples validate on)
+    bootstrap_model="sentence-transformers/all-MiniLM-L6-v2",  # seeds the target geometry
+)
+Trainer(model, TrainingArguments(), train_dataset=rows).train()  # 'mood' auto-detected as an eval-only label
+
+z = model.encode(["a wall of downtuned fuzz that buries the vocals under sheer volume"])  # review snippet -> latent
+print(z.shape)   # (1, 384)
+```
+
+See [`examples/sounds_like/`](examples/sounds_like/) for the full reference task (album review → "how it
+sounds" latent, 481 albums): reproduced through this API at held-out kNN-purity **0.60**, beats-bootstrap
+**4/5**.
+
+## How it works
+
+1. **Bootstrap.** Targets = the bootstrap encoder's embedding of `target_text`. No pairs, no labels.
+2. **Contrastive fit.** InfoNCE (in-batch negatives) trains the LLM emitter to hit its own target and separate
+   from others. A small cosine anchor optionally keeps it tethered to the seed.
+3. **Specialize.** An EMA self-target drifts the geometry off the seed into the model's own arrangement (set
+   `lam_anchor` low to let it go). Emergent structure falls out — never trained, never labeled.
+
+## Validation / early-stop (the part that bites)
+
+Two traps langset refuses to fall into:
+
+- **Never select on training loss** — InfoNCE + EMA can minimize it by *collapsing* the geometry.
+- **Never score retrieval against the frozen bootstrap targets** — the model specializes *away* from them.
+
+So it selects on held-out geometry in the *current* space: input-view↔target-view retrieval + a **collapse
+guard** by default; **held-out kNN-purity (+ beats-bootstrap)** when rows carry geometry labels. Early-stop =
+patience + restore-best.
+
+## Dataset contract
+
+| column | meaning |
+|---|---|
+| `input_text` | what you'll have at inference (a name, query, review) |
+| `target_text` | a description of the same item defining where it lands (seeds the geometry) |
+| *anything else* | optional **eval-only** geometry labels (kNN-purity at validation; never trained on) |
+
+`Trainer` accepts a `datasets.Dataset` or `list[dict]`; use `column_mapping` to rename your columns.
+
+## Using with SetFit
+
+The name is the chain: **lang·set·fit** — a *language* model emits into the *set* geometry (langset, usable on
+its own), which then *fit*s a classifier. `model.as_sentence_transformer()` is a drop-in
+[SetFit](https://github.com/huggingface/setfit) `model_body`, so you can train a few-shot classifier directly
+on the specialized geometry. On the sounds-like example, a genre classifier on the **langset body beats one on
+raw MiniLM, 0.240 vs 0.205**.
+
+### When to reach for which
+
+The clean distinction: **SetFit answers with a *label*; langset answers with a *latent*.** SetFit is a few-shot
+*classifier*; langset is a few-shot *bespoke embedding space* — classification is just one thing you can do
+downstream of a latent.
+
+| | reach for **SetFit** | reach for **langset** |
+|---|---|---|
+| your answer is | a **label** (fixed classes) | a **point in a space** — retrieval, "find similar", ranking, clustering |
+| you define the target by | enumerating classes | a **description** of the geometry ("how it sounds") |
+| your input | text to classify | text *or an identifier* (a name) — leans on the LLM's world knowledge |
+| plain few-shot classification | ✅ simpler, faster, proven | not its job |
+
+- **Use SetFit alone** for plain few-shot classification — it directly optimizes class separation; you won't
+  beat it by bolting on langset.
+- **Use langset** when the answer is a *geometry, not a label* (you'll retrieve / rank / cluster in it).
+- **Use langset → SetFit** when a task-shaped body helps the classifier (demo: genre 0.240 vs 0.205 — modest,
+  so treat it as "plausibly helps", not a slam-dunk).
+
+```bash
+pip install "langset[setfit]"      # pins the verified composition window (below)
+```
+```python
+from sklearn.linear_model import LogisticRegression
+from setfit import SetFitModel
+
+clf = SetFitModel(model_body=langset_model.as_sentence_transformer(),
+                  model_head=LogisticRegression(max_iter=2000),    # direct construction needs an explicit head
+                  labels=[...])
+clf.fit(x_train, y_train, num_epochs=1)        # frozen body + head — the robust path
+clf.predict(["..."])
+```
+
+**Dependency alignment.** SetFit's pins are loose, so versions matter:
+
+| install | transformers / torch | Python | use |
+|---|---|---|---|
+| `langset` | latest (≥4.41) | 3.10+ | modern backbones incl. **Qwen3**; no SetFit |
+| `langset[setfit]` | **4.46.x / <2.5** | **3.10–3.12** | verified SetFit composition |
+
+SetFit imports `transformers.training_args.default_logdir` (removed after 4.46), and 4.46 + torch≥2.5 trips a
+`torch.distributed.tensor` bug — hence the cap. Use the frozen-body `SetFitModel.fit`/`predict` path above; the
+full `setfit.Trainer` (fine-tunes the body) is fragile in this window. Qwen3 + SetFit can't share one env until
+SetFit drops that import.
+
+## Status
+
+v0.1 — the core engine, validated on a real task (album review → "how it sounds" latent), with a downstream
+classifier composition (see above). No trust/hallucination layer yet (intentionally out of v1). Apache-2.0.

langset-0.1.0.dist-info/RECORD

@@ -0,0 +1,11 @@
+langset/__init__.py,sha256=HTWJ-j7EQgzC16xXQwbaSC1iPS6cESVfEYWE24Ilcrs,428
+langset/data.py,sha256=8RZ-_BBRD9J-M3PK4K90i7JgYZrcyhs5ZB22cpK-nL0,1370
+langset/modeling.py,sha256=NoT4smQBBhpC5eO9Cb0UEriXHTC2WKo94c4Da6s3Bq8,8463
+langset/selection.py,sha256=afiZijbATbvIHPNUUU72Kx6TZYFfiXxRJj8FqcUgj-E,3691
+langset/st_module.py,sha256=8wf7xlqh1n8kpN_7AyzcdwfEMiQ24PxqjmBSGeTt2X0,2621
+langset/trainer.py,sha256=Yd6ba1s6mCQLmCtvc0lTBbusMTJ0Jid8F5l_QsAzVQs,6816
+langset/training_args.py,sha256=08Cba4YROFUR-fheWPl-gaPUu358tBJp4liJlILeNE4,1331
+langset-0.1.0.dist-info/METADATA,sha256=MEl_o7lsEbL5aDA9wPn1BZ7pKu3xrTMku5Pw54mIirc,8948
+langset-0.1.0.dist-info/WHEEL,sha256=mffPy8wBnZQn2VnJUU5jE99KsxaSfiyMHV9Yt0aLVxs,87
+langset-0.1.0.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+langset-0.1.0.dist-info/RECORD,,

langset-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.30.1
+Root-Is-Purelib: true
+Tag: py3-none-any

langset-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.