| """miru-codev3-dual: Tokenlearn ⊕ distill_fuse score-fuse via 512-d concat. |
| |
| Cosine on unit concat equals linear score mix: |
| cos(e_q, e_d) = α · cos(A_q, A_d) + (1-α) · cos(B_q, B_d) |
| |
| Default α=0.5 → e = normalize(concat(√α e_A, √(1-α) e_B)). |
| |
| Usage (local folder or Hub repo with tokenlearn/ + distill_fuse/): |
| from dual_encode import DualConcatEncoder |
| enc = DualConcatEncoder.from_pretrained("Takara-DS1/miru-codev3-dual") |
| vecs = enc.encode(["def add(a, b): return a + b"]) |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| from pathlib import Path |
| from typing import Sequence |
|
|
| import numpy as np |
| from huggingface_hub import hf_hub_download, snapshot_download |
| from model2vec import StaticModel |
|
|
|
|
| class DualConcatEncoder: |
| """Two Model2Vec StaticModels fused by √α-scaled concat + L2 norm → 512-d.""" |
|
|
| def __init__( |
| self, |
| model_a: StaticModel, |
| model_b: StaticModel, |
| *, |
| alpha: float = 0.5, |
| max_length: int = 512, |
| ) -> None: |
| if model_a.embedding.shape[1] != model_b.embedding.shape[1]: |
| raise ValueError( |
| f"dim mismatch: {model_a.embedding.shape[1]} vs {model_b.embedding.shape[1]}" |
| ) |
| if not (0.0 <= float(alpha) <= 1.0): |
| raise ValueError("alpha must be in [0, 1]") |
| self.a = model_a |
| self.b = model_b |
| self.alpha = float(alpha) |
| self.max_length = int(max_length) |
| self.dimensions = int(model_a.embedding.shape[1] * 2) |
|
|
| @classmethod |
| def from_pretrained( |
| cls, |
| repo_or_path: str, |
| *, |
| alpha: float | None = None, |
| tokenlearn_subdir: str = "tokenlearn", |
| distill_fuse_subdir: str = "distill_fuse", |
| max_length: int = 512, |
| revision: str | None = None, |
| ) -> DualConcatEncoder: |
| root = Path(repo_or_path) |
| if root.is_dir() and (root / tokenlearn_subdir).is_dir(): |
| base = root |
| else: |
| base = Path( |
| snapshot_download( |
| repo_id=repo_or_path, |
| revision=revision, |
| allow_patterns=[ |
| f"{tokenlearn_subdir}/*", |
| f"{distill_fuse_subdir}/*", |
| "dual_config.json", |
| "README.md", |
| ], |
| ) |
| ) |
|
|
| cfg_path = base / "dual_config.json" |
| cfg: dict = {} |
| if cfg_path.exists(): |
| cfg = json.loads(cfg_path.read_text(encoding="utf-8")) |
| use_alpha = float(alpha if alpha is not None else cfg.get("alpha", 0.5)) |
|
|
| a = StaticModel.from_pretrained(str(base / tokenlearn_subdir)) |
| b = StaticModel.from_pretrained(str(base / distill_fuse_subdir)) |
| return cls(a, b, alpha=use_alpha, max_length=max_length) |
|
|
| @classmethod |
| def from_branches( |
| cls, |
| tokenlearn_id: str, |
| distill_fuse_id: str, |
| *, |
| alpha: float = 0.5, |
| max_length: int = 512, |
| ) -> DualConcatEncoder: |
| """Load the two bags from separate Hub repos / local folders.""" |
| a = StaticModel.from_pretrained(tokenlearn_id) |
| b = StaticModel.from_pretrained(distill_fuse_id) |
| return cls(a, b, alpha=alpha, max_length=max_length) |
|
|
| def fuse(self, ea: np.ndarray, eb: np.ndarray) -> np.ndarray: |
| sa = np.sqrt(self.alpha) |
| sb = np.sqrt(1.0 - self.alpha) |
| out = np.concatenate([sa * ea, sb * eb], axis=-1) |
| out /= np.linalg.norm(out, axis=-1, keepdims=True) + 1e-12 |
| return out.astype(np.float32) |
|
|
| def encode( |
| self, |
| texts: str | Sequence[str], |
| *, |
| max_length: int | None = None, |
| batch_size: int | None = None, |
| ) -> np.ndarray: |
| if isinstance(texts, str): |
| texts = [texts] |
| texts = list(texts) |
| ml = max_length or self.max_length |
| |
| kwargs = {"max_length": ml} |
| if batch_size is not None: |
| kwargs["batch_size"] = batch_size |
| ea = np.asarray(self.a.encode(texts, **kwargs), dtype=np.float32) |
| eb = np.asarray(self.b.encode(texts, **kwargs), dtype=np.float32) |
| if ea.ndim == 1: |
| ea = ea[None, :] |
| eb = eb[None, :] |
| return self.fuse(ea, eb) |
|
|
|
|
| def _smoke() -> None: |
| here = Path(__file__).resolve().parent |
| enc = DualConcatEncoder.from_pretrained(str(here)) |
| v = enc.encode(["def add(a, b):\n return a + b"]) |
| assert v.shape == (1, 512), v.shape |
| assert abs(float(np.linalg.norm(v[0])) - 1.0) < 1e-3 |
| print(f"ok dual encode shape={v.shape} alpha={enc.alpha}") |
|
|
|
|
| if __name__ == "__main__": |
| _smoke() |
|
|
