File size: 607 Bytes
f7fef32 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 |
from sentence_transformers import SentenceTransformer
from torch.nn import EmbeddingBag
import torch
model = SentenceTransformer("tomaarsen/static-retrieval-mrl-en-v1")
embedding_bag: EmbeddingBag = model[0].embedding # type: ignore
embeddings = torch.Tensor(embedding_bag.weight)
assert embeddings.shape == torch.Size([30522, 1024])
print(f"1024 dim - {embeddings.shape[0] * 1024 * 4 / 1024 / 1024:,.1f} MiB:")
print(f"512 dim - {embeddings.shape[0] * 512 * 4 / 1024 / 1024:,.1f} MiB:")
print(f"256 dim - {embeddings.shape[0] * 256 * 4 / 1024 / 1024:,.1f} MiB:")
print("Embeddings[0]", embeddings[0])
|