solailabs/wmt22-cometkiwi-da-pruned-k4

wmt22-cometkiwi-da-pruned-k4

A compressed version of Unbabel/wmt22-cometkiwi-da — a reference-free machine-translation quality estimation model (source + MT only, no human reference required).

Aggressive pruned variant — 4 layers removed. Trades ~8 pt of human Pearson for a smaller model.

What's different from the base model

• —**4 encoder layers dropped (indices 4, 5, 6, 7) out of 24. Layer selection by cosine similarity between each layer's input and output hidden states on a small multilingual calibration set.**
• —layerwise_attention rebuilt to mix only the surviving layers (embeddings + kept layer outputs).
• —No quantization — encoder weights remain fp32.

Accuracy

Benchmarked on 1200 stratified segments from RicardoRei/wmt-da-human-evaluation (reference-free, src+mt only):

All variants at a glance

Usage

Standalone — no gated base-model download. The repo ships everything the loader needs (hparams.yaml + state_dict.pt); the loader instantiates an empty COMET architecture via load_pretrained_weights=False and overlays the fine-tuned weights. Only the ungated microsoft/infoxlm-large tokenizer/config (~5 MB) is fetched on first load and cached.

# pip install "unbabel-comet" "setuptools<81" huggingface_hub pyyaml

from huggingface_hub import snapshot_download
import sys
folder = snapshot_download(repo_id="solailabs/wmt22-cometkiwi-da-pruned-k4")
sys.path.insert(0, folder)
from load import load_model

model = load_model(folder)
out = model.predict(
    [{{"src": "The meeting has been postponed until next week.",
       "mt":  "La réunion a été reportée à la semaine prochaine."}}],
    batch_size=8, gpus=0, progress_bar=False, num_workers=2,
)
print(out["scores"])

No HF_TOKEN required. No license acceptance on Unbabel/wmt22-cometkiwi-da needed.

Files

• —state_dict.pt — model weights (fp32 for -pruned-k2 / -pruned-k4, fp16 for -int8 / -pruned-k4-xs)
• —hparams.yaml — COMET hyper-parameters (encoder model, regressor shape, loss config)
• —config.json — kept/dropped layer indices, quant flag, benchmarked accuracy
• —load.py — drop-in standalone loader
• —README.md — this file

Citation

Base model: `Unbabel/wmt22-cometkiwi-da` by Unbabel.

@inproceedings{{rei-etal-2022-cometkiwi,
    title = "{{C}}omet{{K}}iwi: {{IST}}-{{U}}nbabel 2022 Submission for the Quality Estimation Shared Task",
    author = "Rei, Ricardo  and others",
    booktitle = "WMT 2022",
}}

Released under the same license as the base model (Apache 2.0).