"""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 (in m/s^2, gravity ~9.8) is resampled to a 200-sample, 20 Hz window and placed at one fixed 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. Give the encoder the window the way it was trained: pass your sensor's true sample rate so the window is resampled to 20 Hz and cropped/wrapped to 200 samples (10 s), and supply the accelerometer in m/s^2 (pass unit="g" if your sensor reports g). Both are optional and the API stays backward-compatible: with no sample rate the window is linearly resampled to 200 samples, as before. from inference import TremorEmbedder tr = TremorEmbedder.from_pretrained("EximiusLabs/fusion-embedding-2-tremor") m = tr.embed_motion(accel, sample_rate_hz=50) # accel: np.ndarray [3, T] in m/s^2 scores = tr.rank(accel, ["walking", "sitting", "running", "climbing stairs"], sample_rate_hz=50) Requirements: - torch (CUDA recommended), numpy, transformers>=4.46, huggingface_hub, scipy - 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 GRAVITY = 9.80665 # m/s^2 per g DEFAULT_JOINT = 5 # fixed single joint for a general single-IMU mount, # selected on in-domain held-out data disjoint from # the reported zero-shot sets; supersedes the # checkpoint's legacy 'joint' field 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"] # Placement joint. The general base ('ln' arch) mounts a single IMU at DEFAULT_JOINT, # a fixed skeleton joint selected on in-domain data disjoint from the reported zero-shot # sets. Per-fleet heads ('plain' arch, e.g. Tremor-G1) are trained at, and keep, their # own joint read from the checkpoint. self.joint = self.cfg["joint"] if self.cfg.get("arch") == "plain" else DEFAULT_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, sample_rate_hz: Optional[float] = None, unit: str = "m/s2") -> torch.Tensor: """Raw accelerometer [3, T] -> encoder input [1, 3, win, n_joints, 1] at one joint. Units: the encoder expects m/s^2 (gravity reads ~9.8 at rest), the convention UniMTS was pretrained on. Pass unit="g" if your sensor reports g (gravity ~1.0) and the window is scaled to m/s^2. The default assumes m/s^2 and applies no scaling; unknown input is never silently rescaled. Sampling rate: pass the sensor's true rate (Hz) and the window is resampled to the model's 20 Hz, then wrap-padded or truncated to exactly `win` samples (a 10 s window), matching how the model was trained and evaluated. With no rate the window is linearly resampled to `win` samples (rate-agnostic fallback, backward-compatible).""" a = np.asarray(accel, dtype="float64") if a.ndim != 2 or a.shape[0] != 3: raise ValueError(f"expected accelerometer of shape [3, T], got {list(a.shape)}") if unit == "g": a = a * GRAVITY elif unit not in ("m/s2", "m/s^2"): raise ValueError("unit must be 'm/s2' (default) or 'g'") if sample_rate_hz: from scipy.signal import resample n = max(1, int((a.shape[1] / float(sample_rate_hz)) * 20)) # -> 20 Hz length a = resample(a, n, axis=1) # [3, n] if n < self.win: # wrap-pad short windows a = np.pad(a, ((0, 0), (0, self.win - n)), "wrap") a = a[:, :self.win] # truncate to the 10 s window t = torch.as_tensor(a, dtype=torch.float32).unsqueeze(0) # [1,3,win] else: t = torch.as_tensor(a, dtype=torch.float32).unsqueeze(0) # [1,3,T] if t.shape[-1] != self.win: t = F.interpolate(t, size=self.win, mode="linear", align_corners=False) g = torch.zeros(1, 3, self.win, self.n_joints, 1) g[:, :, :, self.joint, 0] = t return g.to(self.device) @torch.no_grad() def embed_motion(self, accel: np.ndarray, sample_rate_hz: Optional[float] = None, unit: str = "m/s2") -> torch.Tensor: """Embed a raw 3-axis accelerometer window (np.ndarray [3, T], any length). Pass sample_rate_hz (the sensor's true rate) for the trained-time 20 Hz / 200-sample handling, and unit="g" if the sensor reports g rather than m/s^2. See _prep.""" feat = self.enc(self._prep(accel, sample_rate_hz, unit)).squeeze(-1).squeeze(-1).float() 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], sample_rate_hz: Optional[float] = None, unit: str = "m/s2") -> List[tuple]: """Rank candidate activity texts by cosine similarity to a motion window. Pass sample_rate_hz / unit through to embed_motion (see _prep). Returns [(text, score), ...] sorted high to low.""" m = self.embed_motion(accel, sample_rate_hz, unit) 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}")