Views
No views yet
tattabio/gLM2_650M for embedding and retrieval.1import torch
2from transformers import AutoModel, AutoTokenizer
3model = AutoModel.from_pretrained('tattabio/gLM2_650M_embed', torch_dtype=torch.bfloat16, trust_remote_code=True).cuda()
4tokenizer = AutoTokenizer.from_pretrained('tattabio/gLM2_650M_embed', trust_remote_code=True)
5
6# NOTE: Prepend with `<+>` to match gLM2 pre-training.
7sequence = "<+>MALTKVEKRNRIKRRVRGKISGTQASPRLSVYKSNK"
8
9# Tokenize the sequence.
10encodings = tokenizer([sequence], return_tensors='pt')
11# Extract embeddings.
12with torch.no_grad():
13 embeddings = model(encodings.input_ids.cuda()).pooler_output
14
15print(embeddings.shape) # torch.Size([1, 512])