from transformers import AutoTokenizer, AutoModel
import torch

repo_id = "ZoeYou/jepa_lr5e-06_plr0.0001_bs512_tgt512_ema0.996_sigreg0.001_norm_pdim4096_mean_docs5000000"
device = "cuda" if torch.cuda.is_available() else "cpu"

# Load encoder
tokenizer = AutoTokenizer.from_pretrained(repo_id)
model = AutoModel.from_pretrained(repo_id).to(device).eval()

# Mean pooling over non-padding tokens (same logic as project eval script)
def encode(texts, max_length=512):
    inputs = tokenizer(
        texts,
        padding=True,
        truncation=True,
        max_length=max_length,
        return_tensors="pt",
    ).to(device)
    with torch.no_grad():
        out = model(**inputs)
    mask = inputs["attention_mask"].unsqueeze(-1).float()
    emb = (out.last_hidden_state * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1e-12)
    emb = torch.nn.functional.normalize(emb, p=2, dim=1)
    return emb

sep = tokenizer.sep_token
def build_paper_text(title, abstract):
    title = (title or "").strip()
    abstract = (abstract or "").strip()
    return f"{title} {sep} {abstract}" if title and abstract else (title or abstract)

query = build_paper_text(
    "Scientific paper retrieval with domain adaptation",
    "We study representation learning methods for scientific document retrieval.",
)
docs = [
    build_paper_text("Dense retrieval for scientific literature", "A dual-encoder approach for citation recommendation."),
    build_paper_text("Patent semantic search with transformer encoders", "Cross-domain transfer from papers to patents."),
    build_paper_text("Protein folding with large language models", "Large-scale sequence modeling for structure prediction."),
]

q = encode([query])
d = encode(docs)
scores = (d @ q.T).squeeze(-1)
print(scores)

Notes

Max input length is 512 tokens.
Paper text is formatted as title [SEP] abstract.
This repo contains encoder weights only.
For retrieval, L2-normalized embeddings + dot product is recommended.