#!/usr/bin/env python3
"""Minimal embedding example for HRM-Embed-0.6b.

The model is a Hierarchical Reasoning Model (depth-recurrent), NOT a standard transformer
encoder, so it does NOT load via sentence-transformers. Embeddings are the L2-normalized
mean-pool of the final recurrence hidden state (z_h), with bidirectional attention obtained
by passing token_type_ids = attention_mask. Output dim = 1280.
"""
import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModelForCausalLM

MODEL = "viventhraa96/HRM-Embed-0.6b"          # or a local path to the model dir
device = "cuda" if torch.cuda.is_available() else "cpu"

tok = AutoTokenizer.from_pretrained(MODEL)
if tok.pad_token is None:
    tok.pad_token = tok.eos_token
model = AutoModelForCausalLM.from_pretrained(MODEL, trust_remote_code=True, torch_dtype=torch.bfloat16)
model.lm_head = torch.nn.Identity()      # embeddings come from the recurrence state, not the LM head
model = model.to(device).eval()


@torch.no_grad()
def embed(texts, max_length=512):
    tok.padding_side = "right"
    enc = tok(texts, truncation=True, max_length=max_length, padding=True, return_tensors="pt").to(device)
    ids, am = enc["input_ids"], enc["attention_mask"]
    pos = torch.arange(ids.shape[1], device=device).unsqueeze(0).expand(ids.shape[0], -1)
    z_h, _ = model.model(ids, position_ids=pos, use_cache=False, token_type_ids=am)
    mask = am.unsqueeze(-1).to(z_h.dtype)                      # mean-pool over real tokens only
    vec = (z_h * mask).sum(1) / mask.sum(1).clamp_min(1.0)
    return F.normalize(vec.float(), p=2, dim=-1)               # L2-normalized, shape [N, 1280]


if __name__ == "__main__":
    emb = embed([
        "How do I sort a list in Python?",
        "What is the best way to order elements in a Python array?",
        "The mitochondria is the powerhouse of the cell.",
    ])
    print("shape:", tuple(emb.shape))                          # (3, 1280)
    print("cos(similar)  :", f"{float(emb[0] @ emb[1]):.3f}")  # high
    print("cos(unrelated):", f"{float(emb[0] @ emb[2]):.3f}")  # low