how to extract batched embeddings with gLM2?

[seyonechithrananda]

hi @scopello, great work on gLM2 and GAIA, very excited to dive deeper!

is there any specific code for extracting batched embeddings from gLM2 on huggingface?

thanks!

[scopello]

Hi @seyonechithrananda, thanks!

The snippet from huggingface https://huggingface.co/tattabio/gLM2_650M#getting-started can be adapted for batch inference like this:

import torch
from transformers import AutoModel, AutoTokenizer

model = AutoModel.from_pretrained('tattabio/gLM2_650M', torch_dtype=torch.bfloat16, trust_remote_code=True).cuda()
tokenizer = AutoTokenizer.from_pretrained('tattabio/gLM2_650M', trust_remote_code=True)

sequences = ["<+>MALTKVEKRNRIKRRVRGKISGTQASPRLSVYKSNK<+>aatttaaggaa<->MLGIDNIERVKPGGLELVDRLVAVNRVTKVTKGGRAFGFSAIVVVGNED",
             "<+>MALTKVEKRNRIKRRVRGKISGT",
             "<->MLGIDNIERVKPGGLELVDRLVAVNRVTKVTKGGRAFGFSAIVVVGNED",]

encodings = tokenizer(sequences, return_tensors='pt', padding=True)

# Extract batch embeddings.
with torch.no_grad():
    attention_mask = encodings.attention_mask.bool().cuda()
    embeddings = model(encodings.input_ids.cuda(), attention_mask=attention_mask, output_hidden_states=True).last_hidden_state

# Optionally, mean pool.
mask = attention_mask.unsqueeze(-1)
pooled_embeds = torch.where(mask, embeddings, 0.0)
pooled_embeds = torch.sum(pooled_embeds, 1) / torch.sum(mask, dim=1, dtype=pooled_embeds.dtype)

[seyonechithrananda]

Thanks! Do you guys plan to have a more detailed Dataset/dataloader for batches, similar to ESM, to handle padding/batch indexing/etc under the hood for efficient inference? Think it would be nice/allow for people to create larger embedding stores quickly. Right now we wrote some code to batch + pad manually in our SAE codebase, but perhaps this would be best served by a "BatchConverter" class

[seyonechithrananda]

Also, we're playing around with training an SAE on the gLM2_embed activations - would be great to be able to load the "gLM2_650m_embed" version since it's contrastively fine-tuned on AFDB structural clusters and so might pick up interesting structural motifs, but not do pooling by default + extract a variable number of layers from the model for this - would this be possible as a configuration?

[scopello]

Hi @seyonechithrananda, sounds like a cool project! I can work on adding a batched inference script to save embeddings.

Regarding glm2_650M_embed, this model is actually finetuned in 2 stages. The second stage trains a linear layer to align to AFDB clusters, but the backbone is frozen, which is why we did not directly expose the hidden representation. If you still want to extract representations for this model, I can add this or provide the stage 1 finetuning checkpoint.

[ddjamalova]

Hi @scopello , could you please provide the stage 1 finetuning checkpoint? I am also interested in extracting hidden representations.

[scopello]

Hi @ddjamalova, the stage 1 finetuned checkpoint (gLM2 finetuned on Uniref50) can be found here: https://huggingface.co/tattabio/gLM2_650M_UR50
You can follow #7 (comment), and just replace the model name to tattabio/gLM2_650M_UR50