dragonkue/colbert-ko-0.1b

This is a Korean ColBERT model finetuned from skt/A.X-Encoder-base using PyLate. It maps sentences and paragraphs to sequences of 32, 64, 96, or 128-dimensional token embeddings, and uses the MaxSim operator for late-interaction retrieval.

Model Details

Model Description

• Model Type: PyLate model
• Base model: skt/A.X-Encoder-base
• Document Length: 2048 tokens
• Query Length: 32 tokens
• Output Dimensionality: 128 (Matryoshka: 32 / 64 / 96 / 128)
• Similarity Function: MaxSim
• Language: Korean
• License: Apache-2.0

Model Sources

• Documentation: PyLate Documentation
• Repository: PyLate on GitHub
• Hugging Face: PyLate models on Hugging Face

Full Model Architecture

ColBERTWrapper(
  (0): Transformer({'max_seq_length': 2047, 'do_lower_case': False, 'architecture': 'ModernBertModel'})
  (1): Dense({'in_features': 768, 'out_features': 128, 'bias': False, 'activation_function': 'torch.nn.modules.linear.Identity', 'use_residual': False})
)

Usage

This model supports Matryoshka embeddings with multiple dimensions (32, 64, 96, 128) via separate projection heads (Jina-ColBERT-v2 style).

The model ships a bundled model.py so it can be loaded with the transformers library alone — no third-party retrieval package is required. The custom code is opt-in: pass trust_remote_code=True to from_pretrained.

Installation

pip install transformers torch

Quick Start

from transformers import AutoModel, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("dragonkue/colbert-ko-0.1b")
model = AutoModel.from_pretrained(
    "dragonkue/colbert-ko-0.1b",
    trust_remote_code=True,
)
model.eval()
# model.cuda()  # optional

# Pick embedding dimension (32, 64, 96, or 128). Default is 128.
model.set_active_dim(128)

queries = ["검색 쿼리"]
documents = ["문서 내용"]

q_embs = model.encode(tokenizer, queries, is_query=True)
d_embs = model.encode(tokenizer, documents, is_query=False)

print(f"Query shape: {q_embs[0].shape}")  # (num_tokens, 128)
print(f"Doc shape:   {d_embs[0].shape}")  # (num_tokens, 128)

# ColBERT late-interaction (MaxSim) score for one (query, doc) pair
score = model.maxsim(q_embs[0], d_embs[0])
print(f"MaxSim: {score:.4f}")

Reranking

query = "대한민국의 수도는 어디인가요?"
candidates = [
    "대한민국의 수도는 서울이며, 인구가 가장 많은 도시이다.",
    "파리는 프랑스의 수도이고 에펠탑으로 유명하다.",
    "오늘 날씨는 맑고 기온은 25도 정도이다.",
]

q_embs = model.encode(tokenizer, [query], is_query=True)
c_embs = model.encode(tokenizer, candidates, is_query=False)

ranked = sorted(
    [(i, model.maxsim(q_embs[0], c_embs[i])) for i in range(len(candidates))],
    key=lambda x: x[1],
    reverse=True,
)
for rank, (i, score) in enumerate(ranked, start=1):
    print(f"  rank {rank}: score={score:.4f}  | {candidates[i]}")

Approximate Retrieval (PyLate PLAID, optional)

For PLAID indexing over a large corpus you can additionally use PyLate. PyLate loads the model as a standard sentence-transformers checkpoint (no trust_remote_code needed). The 128-dim head saved at 1_Dense/model.safetensors matches the Matryoshka 128-dim head, so PyLate produces the same embeddings as the snippet above.

from pylate import indexes, models, retrieve

model = models.ColBERT(model_name_or_path="dragonkue/colbert-ko-0.1b")

index = indexes.PLAID(
    index_folder="pylate-index",
    index_name="index",
    override=True,
)

documents_ids = ["1", "2", "3"]
documents = ["첫번째 문서입니다", "두번째 문서입니다", "세번째 문서입니다"]

documents_embeddings = model.encode(documents, is_query=False)
index.add_documents(
    documents_ids=documents_ids,
    documents_embeddings=documents_embeddings,
)

retriever = retrieve.ColBERT(index=index)
queries_embeddings = model.encode(["첫번째 문서 검색"], is_query=True)
scores = retriever.retrieve(queries_embeddings=queries_embeddings, k=3)
print(scores)

The PyLate path is fixed to the 128-dim head; for Matryoshka dimension switching (32 / 64 / 96 / 128), use the trust_remote_code=True path shown above.

Evaluation Results (NDCG@10)

Comparison with Other Models (dim128)

Performance by Embedding Dimension

Differences across dimensions are small (within ≈0.01 on each benchmark) and can flip between runs on the more saturated benchmarks such as AutoRAG. The 128-dim head is the default; smaller dimensions trade a few points of NDCG for a 2–4× reduction in vector size.

Scores for dragonkue-colbert-ko-0.1b are measured here with raw MaxSim (no PLAID quantization) on transformers==5.3.0 and pytorch==2.9. Scores for the other models are reproduced from prior reports and may differ from a re-run on the same stack.

• Loss: src.losses.MatryoshkaColBERTLoss with these parameters:

  {
      "dims": [
          32,
          64,
          96,
          128
      ],
      "weights": [
          0.25,
          0.25,
          0.25,
          0.25
      ],
      "temperature": 1.0
  }
  

Training Hyperparameters

Non-Default Hyperparameters

• eval_strategy: steps
• per_device_train_batch_size: 64
• learning_rate: 1e-05
• num_train_epochs: 2
• warmup_ratio: 0.1
• fp16: True
• dataloader_drop_last: True
• gradient_checkpointing: True
• gradient_checkpointing_kwargs: {'use_reentrant': False}
• router_mapping: {'anchor': 'query', 'positive': 'document', 'neg_0': 'document'}

All Hyperparameters

• overwrite_output_dir: False
• do_predict: False
• eval_strategy: steps
• prediction_loss_only: True
• per_device_train_batch_size: 64
• per_device_eval_batch_size: 8
• per_gpu_train_batch_size: None
• per_gpu_eval_batch_size: None
• gradient_accumulation_steps: 1
• eval_accumulation_steps: None
• torch_empty_cache_steps: None
• learning_rate: 1e-05
• weight_decay: 0.0
• adam_beta1: 0.9
• adam_beta2: 0.999
• adam_epsilon: 1e-08
• max_grad_norm: 1.0
• num_train_epochs: 2
• max_steps: -1
• lr_scheduler_type: linear
• lr_scheduler_kwargs: {}
• warmup_ratio: 0.1
• warmup_steps: 0
• log_level: passive
• log_level_replica: warning
• log_on_each_node: True
• logging_nan_inf_filter: True
• save_safetensors: True
• save_on_each_node: False
• save_only_model: False
• restore_callback_states_from_checkpoint: False
• no_cuda: False
• use_cpu: False
• use_mps_device: False
• seed: 42
• data_seed: None
• jit_mode_eval: False
• use_ipex: False
• bf16: False
• fp16: True
• fp16_opt_level: O1
• half_precision_backend: auto
• bf16_full_eval: False
• fp16_full_eval: False
• tf32: None
• local_rank: 0
• ddp_backend: None
• tpu_num_cores: None
• tpu_metrics_debug: False
• debug: []
• dataloader_drop_last: True
• dataloader_num_workers: 0
• dataloader_prefetch_factor: None
• past_index: -1
• disable_tqdm: False
• remove_unused_columns: True
• label_names: None
• load_best_model_at_end: False
• ignore_data_skip: False
• fsdp: []
• fsdp_min_num_params: 0
• fsdp_config: {'min_num_params': 0, 'xla': False, 'xla_fsdp_v2': False, 'xla_fsdp_grad_ckpt': False}
• fsdp_transformer_layer_cls_to_wrap: None
• accelerator_config: {'split_batches': False, 'dispatch_batches': None, 'even_batches': True, 'use_seedable_sampler': True, 'non_blocking': False, 'gradient_accumulation_kwargs': None}
• parallelism_config: None
• deepspeed: None
• label_smoothing_factor: 0.0
• optim: adamw_torch_fused
• optim_args: None
• adafactor: False
• group_by_length: False
• length_column_name: length
• ddp_find_unused_parameters: None
• ddp_bucket_cap_mb: None
• ddp_broadcast_buffers: False
• dataloader_pin_memory: True
• dataloader_persistent_workers: False
• skip_memory_metrics: True
• use_legacy_prediction_loop: False
• push_to_hub: False
• resume_from_checkpoint: None
• hub_model_id: None
• hub_strategy: every_save
• hub_private_repo: None
• hub_always_push: False
• hub_revision: None
• gradient_checkpointing: True
• gradient_checkpointing_kwargs: {'use_reentrant': False}
• include_inputs_for_metrics: False
• include_for_metrics: []
• eval_do_concat_batches: True
• fp16_backend: auto
• push_to_hub_model_id: None
• push_to_hub_organization: None
• mp_parameters:
• auto_find_batch_size: False
• full_determinism: False
• torchdynamo: None
• ray_scope: last
• ddp_timeout: 1800
• torch_compile: False
• torch_compile_backend: None
• torch_compile_mode: None
• include_tokens_per_second: False
• include_num_input_tokens_seen: False
• neftune_noise_alpha: None
• optim_target_modules: None
• batch_eval_metrics: False
• eval_on_start: False
• use_liger_kernel: False
• liger_kernel_config: None
• eval_use_gather_object: False
• average_tokens_across_devices: False
• prompts: None
• batch_sampler: batch_sampler
• multi_dataset_batch_sampler: proportional
• router_mapping: {'anchor': 'query', 'positive': 'document', 'neg_0': 'document'}
• learning_rate_mapping: {}

Training Logs

Framework Versions

• Python: 3.11.13
• Sentence Transformers: 5.1.1
• PyLate: 1.3.4
• Transformers: 4.56.2
• PyTorch: 2.8.0+cu128
• Accelerate: 1.12.0
• Datasets: 4.4.1
• Tokenizers: 0.22.2-rc0

Citation

BibTeX

Sentence Transformers

@inproceedings{reimers-2019-sentence-bert,
    title = "Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks",
    author = "Reimers, Nils and Gurevych, Iryna",
    booktitle = "Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing",
    month = "11",
    year = "2019",
    publisher = "Association for Computational Linguistics",
    url = "https://arxiv.org/abs/1908.10084"
}

PyLate

@misc{PyLate,
title={PyLate: Flexible Training and Retrieval for Late Interaction Models},
author={Chaffin, Antoine and Sourty, Raphaël},
url={https://github.com/lightonai/pylate},
year={2024}
}

Jina ColBERT v2

@article{jina-colbert-v2,
    title={Jina-ColBERT-v2: A General-Purpose Multilingual Late Interaction Retriever},
    author={Rohan Jha and Bo Wang and Michael Günther and Saba Sturua and Mohammad Kalim Akram and Han Xiao},
    year={2024},
    journal={arXiv preprint arXiv:2408.16672},
    url={https://arxiv.org/abs/2408.16672}
}