Spaces:
Running
Running
File size: 2,348 Bytes
88d2f2a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 | """Trigger the first COMET model download + score one (src, mt) pair.
Run with the project venv:
.venv/bin/python scripts/test_comet.py
"""
from __future__ import annotations
import time
from comet import download_model, load_from_checkpoint
# Per spec we prefer Unbabel/wmt22-cometkiwi-da (reference-free), but that
# model is a *gated* HF repo and the locally-stored HF token lacks
# ``canReadGatedRepos`` permission. Fall back to the non-gated WMT20 QE
# DA-v2 model (also reference-free) so the judge can run end-to-end.
PREFERRED_MODEL = "Unbabel/wmt22-cometkiwi-da"
FALLBACK_MODEL = "Unbabel/wmt20-comet-qe-da"
def main() -> None:
MODEL_ID = PREFERRED_MODEL
print(f"[test_comet] trying preferred {MODEL_ID} (gated HF)...")
t0 = time.perf_counter()
try:
ckpt_path = download_model(MODEL_ID)
except Exception as exc: # noqa: BLE001 - we want broad fallback
print(f"[test_comet] preferred failed: {type(exc).__name__}: {exc}")
MODEL_ID = FALLBACK_MODEL
print(f"[test_comet] falling back to {MODEL_ID} (non-gated)...")
t0 = time.perf_counter()
ckpt_path = download_model(MODEL_ID)
print(f"[test_comet] checkpoint at: {ckpt_path}")
print(f"[test_comet] download/locate took {time.perf_counter() - t0:.2f}s")
print("[test_comet] loading model from checkpoint...")
t1 = time.perf_counter()
model = load_from_checkpoint(ckpt_path)
print(f"[test_comet] load took {time.perf_counter() - t1:.2f}s")
data = [
{
"src": "央行行长潘功胜在金融街论坛年会上表示，将根据需要适时降准。",
"mt": (
"Will the People's Bank of China (PBOC) announce a cut to the "
"Reserve Requirement Ratio (RRR) before August 23, 2026?"
),
}
]
print("[test_comet] running predict on 1 pair (CPU)...")
t2 = time.perf_counter()
output = model.predict(data, batch_size=1, progress_bar=False, gpus=0)
elapsed = time.perf_counter() - t2
system_score = getattr(output, "system_score", None)
scores = getattr(output, "scores", None)
print(f"[test_comet] system_score = {system_score}")
print(f"[test_comet] scores = {scores}")
print(f"[test_comet] predict took {elapsed:.2f}s")
if __name__ == "__main__":
main()
|