Views
No views yet
macwiatrak/baclm-350m-masked is a 350M-parameter masked language model for bacterial genomics. It is designed to model both protein sequences and intergenic DNA with a single shared character-level transformer encoder.token_type_ids, which let the model distinguish modalities internally. Protein and DNA examples can be batched together, but each example should correspond to a single sequence modality.MKTAYIAKQRQISFVKSHFSRQatgcttagctagcttacg1import torch
2from transformers import AutoModel, AutoTokenizer
3
4model_name = "macwiatrak/baclm-350m-masked"
5
6tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
7model = AutoModel.from_pretrained(model_name, trust_remote_code=True, dtype=torch.bfloat16)
8model.eval().cuda()
9
10seqs = [
11 "MKTAYIAKQRQISFVKSHFSRQ", # protein: uppercase
12 "atgcttagctagcttacg", # DNA: lowercase
13]
14
15batch = tokenizer.batch_encode_plus(
16 seqs,
17 padding=True,
18 truncation=True,
19 max_length=2048,
20 return_tensors="pt",
21)
22batch = {k: v.cuda() for k, v in batch.items()}
23
24with torch.no_grad():
25 outputs = model(
26 input_ids=batch["input_ids"],
27 token_type_ids=batch.get("token_type_ids"),
28 attention_mask=batch.get("attention_mask"),
29 )
30
31# Token-level embeddings
32token_embeddings = outputs.last_hidden_state
33
34# Mean pooled embeddings
35attention_mask = batch["attention_mask"].unsqueeze(-1)
36mean_embeddings = (token_embeddings * attention_mask).sum(dim=1) / attention_mask.sum(dim=1).clamp_min(1)
37print(mean_embeddings.shape)TBD