"""Tremor v0.1 inference — search body-worn motion in natural language. Tremor maps a short window of accelerometer motion into the Qwen3-VL-Embedding-2B text space, so a stream of inertial motion becomes retrievable with plain-language activity queries ("walking upstairs", "sitting", "picking something up"). Pipeline: a raw 3-axis accelerometer window is resampled to 200 samples at 20 Hz and placed at one joint of a frozen UniMTS ST-GCN encoder, which outputs a 512-d motion feature; a small trained projector (this repository's only weights) maps that to the frozen Qwen base's 2048-d text space. Motion and text embeddings are L2-normalized and compared by cosine. from inference import TremorEmbedder tr = TremorEmbedder.from_pretrained("EximiusLabs/fusion-embedding-2-tremor") m = tr.embed_motion(accel) # accel: np.ndarray [3, T], any length scores = tr.rank(accel, ["walking", "sitting", "running", "climbing stairs"]) Requirements: - torch (CUDA recommended), numpy, transformers>=4.46, huggingface_hub - The frozen UniMTS encoder code and weights (Apache-2.0): git clone https://github.com/xiyuanzh/UniMTS # provides model.py::ST_GCN_18 Point UNIMTS_REPO at the clone, or pass unimts_repo= to from_pretrained. The UniMTS weights (checkpoint/UniMTS.pth) download automatically from the HF hub. The frozen Qwen3-VL-Embedding-2B base downloads from its original repository. Embedding quality is sensitive to the base's chat-template formatting; use the template provided here rather than constructing your own. """ from __future__ import annotations import json import os import sys from typing import List, Optional, Sequence, Union import numpy as np import torch import torch.nn as nn import torch.nn.functional as F BASE_MODEL = "Qwen/Qwen3-VL-Embedding-2B" UNIMTS_REPO = os.environ.get("UNIMTS_REPO", "UniMTS") # local clone of xiyuanzh/UniMTS CKPT_FILE = "tremor_projector.pt" SAFETENSORS_FILE = "model.safetensors" # projector weights; config in the file's metadata def _chat(text: str) -> str: """The Qwen base's embedding chat-template. Motion is matched against text embedded this way.""" return ("<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n" f"<|im_start|>user\n{text}<|im_end|>\n<|im_start|>assistant\n") def _build_projector(in_dim: int, out_dim: int, arch: str = "ln") -> nn.Sequential: """Projector head. 'ln' = the general base (LayerNorm front); 'plain' = the per-fleet variants (e.g. Tremor-G1), which omit the input LayerNorm. Selected by the checkpoint config.""" if arch == "plain": return nn.Sequential(nn.Linear(in_dim, 1024), nn.GELU(), nn.Dropout(0.3), nn.Linear(1024, out_dim)) return nn.Sequential(nn.LayerNorm(in_dim), nn.Linear(in_dim, 1024), nn.GELU(), nn.Dropout(0.3), nn.Linear(1024, out_dim)) class TremorEmbedder: def __init__(self, ckpt_path: str, device: str = "cuda", dtype=torch.bfloat16, unimts_repo: str = UNIMTS_REPO): from huggingface_hub import hf_hub_download from transformers import AutoModel, AutoTokenizer self.device = device if ckpt_path.endswith(".safetensors"): import safetensors.torch as st from safetensors import safe_open proj_sd = st.load_file(ckpt_path) with safe_open(ckpt_path, framework="pt") as f: self.cfg = json.loads((f.metadata() or {})["config"]) else: # legacy .pt checkpoint ck = torch.load(ckpt_path, map_location="cpu", weights_only=False) self.cfg = ck["config"]; proj_sd = ck["proj"] self.joint = self.cfg["joint"] self.n_joints = self.cfg["n_joints"] self.win = self.cfg["window_samples"] # frozen UniMTS motion encoder (accelerometer-only ST-GCN) if not os.path.isdir(unimts_repo): raise FileNotFoundError( f"UniMTS repo not found at '{unimts_repo}'. Clone it and set unimts_repo=/UNIMTS_REPO:\n" " git clone https://github.com/xiyuanzh/UniMTS") sys.path.insert(0, unimts_repo) from model import ST_GCN_18 w = hf_hub_download(self.cfg["unimts_repo"], self.cfg["unimts_file"]) sd = torch.load(w, map_location="cpu", weights_only=False) sd = sd.get("state_dict", sd) acc = {k[len("acc."):]: v for k, v in sd.items() if k.startswith("acc.")} enc = ST_GCN_18(in_channels=3) enc.load_state_dict(acc, strict=False) self.enc = enc.eval().to(device) for p in self.enc.parameters(): p.requires_grad_(False) # trained projector (this repository) self.proj = _build_projector(self.cfg["in_dim"], self.cfg["out_dim"], self.cfg.get("arch", "ln")).to(device) self.proj.load_state_dict({k: v.float() for k, v in proj_sd.items()}) self.proj.eval() # frozen Qwen text space self.base = AutoModel.from_pretrained(BASE_MODEL, trust_remote_code=True, dtype=dtype).to(device).eval() for p in self.base.parameters(): p.requires_grad_(False) self.tok = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True) self.tok.padding_side = "right" @classmethod def from_pretrained(cls, repo_or_path: str, device: str = "cuda", revision: Optional[str] = None, unimts_repo: str = UNIMTS_REPO, subfolder: str = "", **kw) -> "TremorEmbedder": """Load from a local checkpoint (tremor_projector.pt) or an HF repo. subfolder selects a variant that ships alongside the general base in the same repo, e.g. the Unitree-specialized head: from_pretrained("EximiusLabs/fusion-embedding-2-tremor", subfolder="g1").""" if os.path.isfile(repo_or_path): path = repo_or_path elif os.path.isdir(repo_or_path): cand = os.path.join(repo_or_path, subfolder, SAFETENSORS_FILE) path = cand if os.path.exists(cand) else os.path.join(repo_or_path, subfolder, CKPT_FILE) else: from huggingface_hub import hf_hub_download from huggingface_hub.utils import EntryNotFoundError cfg_fn = f"{subfolder}/config.json" if subfolder else "config.json" try: # fetch config.json so a real load registers a Hub download hf_hub_download(repo_or_path, cfg_fn, revision=revision) except Exception: pass st_fn = f"{subfolder}/{SAFETENSORS_FILE}" if subfolder else SAFETENSORS_FILE try: # prefer safetensors, fall back to the legacy pickle path = hf_hub_download(repo_or_path, st_fn, revision=revision) except EntryNotFoundError: pt_fn = f"{subfolder}/{CKPT_FILE}" if subfolder else CKPT_FILE path = hf_hub_download(repo_or_path, pt_fn, revision=revision) return cls(path, device=device, unimts_repo=unimts_repo, **kw) # --------------------------------------------------------------- motion def _prep(self, accel: np.ndarray) -> torch.Tensor: """Raw accelerometer [3, T] -> encoder input [1, 3, win, n_joints, 1] at one joint.""" a = torch.as_tensor(np.asarray(accel), dtype=torch.float32) if a.ndim != 2 or a.shape[0] != 3: raise ValueError(f"expected accelerometer of shape [3, T], got {list(a.shape)}") a = a.unsqueeze(0) # [1,3,T] if a.shape[-1] != self.win: # resample time axis to the model rate a = F.interpolate(a, size=self.win, mode="linear", align_corners=False) g = torch.zeros(1, 3, self.win, self.n_joints, 1) g[:, :, :, self.joint, 0] = a return g.to(self.device) @torch.no_grad() def embed_motion(self, accel: np.ndarray) -> torch.Tensor: """Embed a raw 3-axis accelerometer window (np.ndarray [3, T], any length, any rate).""" feat = self.enc(self._prep(accel)).squeeze(-1).squeeze(-1).float() # [1,512] return F.normalize(self.proj(feat), dim=-1).squeeze(0).cpu() # --------------------------------------------------------------- text @torch.no_grad() def embed_text(self, text: Union[str, Sequence[str]]) -> torch.Tensor: one = isinstance(text, str) texts = [text] if one else list(text) out = [] for t in texts: # padding-free (base is sensitive to padding) enc = self.tok(_chat(t), return_tensors="pt", truncation=True, max_length=64).to(self.device) h = self.base(**enc).last_hidden_state idx = int(enc["attention_mask"].sum().item()) - 1 out.append(F.normalize(h[0, idx].float(), dim=-1).cpu()) e = torch.stack(out) return e.squeeze(0) if one else e # --------------------------------------------------------------- readout @torch.no_grad() def rank(self, accel: np.ndarray, texts: Sequence[str]) -> List[tuple]: """Rank candidate activity texts by cosine similarity to a motion window. Returns [(text, score), ...] sorted high to low.""" m = self.embed_motion(accel) te = self.embed_text(list(texts)) scores = (te @ m).tolist() return sorted(zip(texts, scores), key=lambda x: -x[1]) if __name__ == "__main__": # smoke: random motion, four candidate activities (needs the models + UniMTS clone to run) tr = TremorEmbedder.from_pretrained(os.path.join(os.path.dirname(__file__), "out")) demo = np.random.randn(3, 300).astype("float32") for text, score in tr.rank(demo, ["walking", "sitting", "running", "climbing stairs"]): print(f"{score:+.3f} {text}")