chengCCC/CrossDNA_pretrain

Using CrossDNA 8.1M

import os
os.environ.setdefault("DISABLE_TORCH_COMPILE", "1")  

import torch
if hasattr(torch, "compile"):
    def _no_compile(fn=None, *args, **kwargs):
        if fn is None:
            def deco(f): return f
            return deco
        return fn
    torch.compile = _no_compile

import torch
from transformers import AutoTokenizer, AutoModelForMaskedLM

# Hugging Face Remote Repository Example
# repo_id = "chengCCC/CrossDNA_pretrain"
# subdir = "8.1M"

# tok = AutoTokenizer.from_pretrained(repo_id, subfolder=subdir, trust_remote_code=True, local_files_only=True)
# model = AutoModelForMaskedLM.from_pretrained(repo_id, subfolder=subdir, trust_remote_code=True, local_files_only=True).eval()

# Local Model Example
MODEL_DIR = "/data/zhaol/projects/huggingface_crossdna/crossdna"
tok = AutoTokenizer.from_pretrained(MODEL_DIR, trust_remote_code=True, local_files_only=True)
model = AutoModelForMaskedLM.from_pretrained(MODEL_DIR, trust_remote_code=True, local_files_only=True).eval()

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# 512bp DNA sequence
seq = "ACGT" * 128


enc = tok(seq, return_tensors="pt", add_special_tokens=False)


x = enc["input_ids"].to(device)  # [1, L]

# Key: map tokenizer IDs back to the model-required range [0..4]
#   A, C, G, T, N: 7..11 -> 0..4
#   All other tokens (UNK/PAD/SEP/CLS/...) are treated as N (=4)
x = torch.where(x >= 7, x - 7, torch.full_like(x, 4))

# ======  embedding ======
was_pretrain = getattr(model.backbone, "pretrain", False)
was_for_repr = getattr(model.backbone, "for_representation", False)
model.backbone.pretrain = False
model.backbone.for_representation = True

with torch.inference_mode():
    embeddings, _ = model.backbone(x)   # [B, L, H]

print("input_ids.shape =", tuple(x.shape)) # input_ids.shape = (1, 512)
print("embeddings.shape =", tuple(embeddings.shape)) # embeddings.shape = (1, 512, 128)

model.backbone.pretrain = was_pretrain
model.backbone.for_representation = was_for_repr