geolip-esm2_t33_650M_UR50D

from huggingface_hub import hf_hub_download
from geolip_core.pipeline.esm2_geometric import ESM2GeometricPipeline
import torch

# Load
pipe = ESM2GeometricPipeline('esm2_geo')
ckpt = torch.load(
    hf_hub_download('AbstractPhil/geolip-esm2_t33_650M_UR50D',
                    'prototype/v1_distill/epoch_6.pt'),
    map_location='cuda')
current = pipe.state_dict()
current.update({k: v for k, v in ckpt['state_dict'].items() if k in current})
pipe.load_state_dict(current, strict=False)
pipe.eval().cuda()

# Run
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained('facebook/esm2_t33_650M_UR50D')
enc = tokenizer("MVLSPADKTNVKAAWGKVGAHAGEYGAEALERMFLSFPTTKTYFPHFDLSH",
                return_tensors='pt', padding=True).to('cuda')

with torch.no_grad():
    loss, info = pipe.forward_distill(enc['input_ids'], enc['attention_mask'])

# Diagnostics
print(info['gate_info'])
print(info['svd_S'])
print(pipe.cache_get('embedding'))  # protein-level geometric feature