Fix the dense embedding dimension slice
#32
by CyberNativeAI - opened
The dimension option in the dense Transformers example currently slices the batch axis after selecting the first token. This one-line change slices the feature axis before the existing normalization.
With dimension=128, four inputs still produce (4, 768); 129 inputs produce (128, 768), dropping the last input. The corrected expression produces (4, 128) and (129, 128). The example's default dimension=768 is unchanged for its four-input batch.
Minimal reproduction, without model downloads:
from types import SimpleNamespace
import torch
for batch, dimension in ((4, 128), (129, 128), (4, 256), (4, 768)):
outputs = SimpleNamespace(
last_hidden_state=torch.arange(batch * 2 * 768).reshape(batch, 2, 768)
)
old = outputs.last_hidden_state[:, 0][:dimension]
new = outputs.last_hidden_state[:, 0, :dimension]
assert new.shape == (batch, dimension)
print(batch, dimension, tuple(old.shape), "->", tuple(new.shape))
Checked with PyTorch 2.4.0 on CPU. Additional synthetic assertions verified retained values, every input row and normalization after slicing. Model loading, actual embeddings and retrieval quality were not tested.
Submitted by CyberNative AI LLC; AI-assisted.