cosdoc-nq2

CosDoc — Attention nq=2

CosDoc is a visual document embedding model trained with supervised metric learning and hard-example selection via a Reinforcement Learning professor network.

Pooler variant: Attention nq=2 — Multi-query attention pooler, num_queries=2.

Architecture

Performance (LA-CDIP, full validation pairs)

Source run: Sprint3b_S0_subcenter_cosface_seed42_noinit_nq2_fase1_E10

Usage

import torch
from huggingface_hub import hf_hub_download
from cavl_doc.models.backbone_loader import load_model
from cavl_doc.models.modeling_cavl import build_cavl_model

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

# Download fine-tuned weights
ckpt_path = hf_hub_download(repo_id="Jpcosta90/cosdoc-nq2", filename="best_model.pt")
ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
cfg  = ckpt["config"]

backbone, _, tokenizer, _, _ = load_model("InternVL3-2B")
model = build_cavl_model(
    backbone=backbone,
    cut_layer=cfg["cut_layer"],
    pooler_type=cfg["pooler_type"],
    num_queries=cfg.get("num_queries", 1),
)
model.pool.load_state_dict(ckpt["siam_pool"])
model.head.load_state_dict(ckpt["siam_head"])
model.eval().to(device)

Citation

@misc{cosdoc2026,
  title  = {CosDoc: Cosine-Margin Document Embeddings with RL-guided Hard Mining},
  author = {Costa, João Paulo},
  year   = {2026},
  url    = {https://huggingface.co/Jpcosta90/cosdoc-nq2}
}