import torch
from transformers import AutoTokenizer, AutoModel

model_id = "heyongxin233/DETree"
tgt_layer = 18  # 0=embeddings; 1..24=encoder layers (RoBERTa-large)
device = "cuda" if torch.cuda.is_available() else "cpu"
print("Using device:", device)

tok = AutoTokenizer.from_pretrained(model_id)
enc = AutoModel.from_pretrained(model_id, output_hidden_states=True).to(device)
enc.eval()

texts = ["An example sentence.", "Another one."]
batch = tok(texts, padding=True, truncation=True, return_tensors="pt")
batch = {k: v.to(device) for k, v in batch.items()}

with torch.inference_mode():
    out = enc(**batch)
    hs = out.hidden_states[tgt_layer]            # (bsz, seq, hidden)
    mask = batch["attention_mask"].unsqueeze(-1)  # (bsz, seq, 1)
    hs = hs.masked_fill(~mask.bool(), float("-inf"))
    emb, _ = hs.max(dim=1)                       # max-pool over tokens
    emb = torch.nn.functional.normalize(emb, p=2, dim=-1)

print(emb.device, emb.shape)  # -> (batch_size, 1024)