---
datasets:
- tattabio/OMG
license: apache-2.0
---
# gLM2_650M_embed

gLM2_embed is a fine-tuned vesion of [`tattabio/gLM2_650M`](https://huggingface.co/tattabio/gLM2_650M) for embedding and retrieval.

- The first stage finetunes gLM2 over one epoch of UniRef50.
- The second stage trains an adapter layer to align mean-pooled representations with AlphaFold structural [clusters](https://www.nature.com/articles/s41586-023-06510-w).

## Getting Started
```python
import torch
from transformers import AutoModel, AutoTokenizer
model = AutoModel.from_pretrained('tattabio/gLM2_650M_embed', torch_dtype=torch.bfloat16, trust_remote_code=True).cuda()
tokenizer = AutoTokenizer.from_pretrained('tattabio/gLM2_650M_embed', trust_remote_code=True)

# NOTE: Prepend with `<+>` to match gLM2 pre-training.
sequence = "<+>MALTKVEKRNRIKRRVRGKISGTQASPRLSVYKSNK"

# Tokenize the sequence.
encodings = tokenizer([sequence], return_tensors='pt')
# Extract embeddings.
with torch.no_grad():
    embeddings = model(encodings.input_ids.cuda()).pooler_output

print(embeddings.shape)  # torch.Size([1, 512])
```