|
| 1 | +"""CommonLingua: PleIAs' byte-level LID model (PleIAs/CommonLingua). |
| 2 | +
|
| 3 | +Requires the ``commonlid[commonlingua]`` extra (torch). The checkpoint |
| 4 | +embeds its own ``lang2idx`` map, so no separate metadata file is fetched. |
| 5 | +Device selection mirrors AfroLID: MPS > CUDA > CPU. |
| 6 | +""" |
| 7 | + |
| 8 | +from __future__ import annotations |
| 9 | + |
| 10 | +from collections.abc import Sequence |
| 11 | +from typing import TYPE_CHECKING, Any, ClassVar |
| 12 | + |
| 13 | +from commonlid.core.lid_model import LIDModel |
| 14 | +from commonlid.core.registry import register_model |
| 15 | + |
| 16 | +if TYPE_CHECKING: |
| 17 | + import torch |
| 18 | + |
| 19 | + |
| 20 | +@register_model |
| 21 | +class CommonLinguaModel(LIDModel): |
| 22 | + model_id = "commonlingua" |
| 23 | + # Byte-level model: casing carries strong language signal, so we feed |
| 24 | + # raw UTF-8 and skip the OpenLID normer (which lowercases everything). |
| 25 | + requires_preprocessing: ClassVar[bool] = False |
| 26 | + |
| 27 | + _REPO_ID: ClassVar[str] = "PleIAs/CommonLingua" |
| 28 | + _CHECKPOINT_FILENAME: ClassVar[str] = "model.pt" |
| 29 | + _INTERNAL_BATCH: ClassVar[int] = 256 |
| 30 | + |
| 31 | + def __init__(self) -> None: |
| 32 | + super().__init__() |
| 33 | + self._model: Any = None |
| 34 | + self._idx2lang: dict[int, str] | None = None |
| 35 | + self._max_len: int | None = None |
| 36 | + self._device: str | None = None |
| 37 | + |
| 38 | + def load(self) -> None: |
| 39 | + if self._loaded: |
| 40 | + return |
| 41 | + try: |
| 42 | + import torch |
| 43 | + except ImportError as exc: |
| 44 | + msg = "CommonLingua requires torch. Install with: pip install 'commonlid[commonlingua]'" |
| 45 | + raise ImportError(msg) from exc |
| 46 | + |
| 47 | + from huggingface_hub import hf_hub_download |
| 48 | + |
| 49 | + from commonlid.vendor.commonlingua.model import CONFIGS, ByteHybrid |
| 50 | + |
| 51 | + ckpt_path = hf_hub_download(repo_id=self._REPO_ID, filename=self._CHECKPOINT_FILENAME) |
| 52 | + ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False) |
| 53 | + |
| 54 | + if torch.backends.mps.is_available(): |
| 55 | + device = "mps" |
| 56 | + elif torch.cuda.is_available(): |
| 57 | + device = "cuda" |
| 58 | + else: |
| 59 | + device = "cpu" |
| 60 | + |
| 61 | + model = ByteHybrid( # type: ignore[no-untyped-call] |
| 62 | + num_classes=ckpt["num_classes"], |
| 63 | + max_len=ckpt["max_len"], |
| 64 | + **CONFIGS[ckpt["config"]], |
| 65 | + ) |
| 66 | + model.load_state_dict(ckpt["model_state_dict"]) |
| 67 | + model.eval().to(device) |
| 68 | + |
| 69 | + self._model = model |
| 70 | + self._idx2lang = {v: k for k, v in ckpt["lang2idx"].items()} |
| 71 | + self._max_len = int(ckpt["max_len"]) |
| 72 | + self._device = device |
| 73 | + super().load() |
| 74 | + |
| 75 | + def _encode(self, texts: Sequence[str]) -> torch.Tensor: |
| 76 | + import numpy as np |
| 77 | + import torch |
| 78 | + |
| 79 | + assert self._max_len is not None |
| 80 | + out = np.full((len(texts), self._max_len), 256, dtype=np.int64) |
| 81 | + for i, t in enumerate(texts): |
| 82 | + raw = t.encode("utf-8", errors="replace")[: self._max_len] |
| 83 | + if raw: |
| 84 | + out[i, : len(raw)] = np.frombuffer(raw, dtype=np.uint8) |
| 85 | + return torch.from_numpy(out) |
| 86 | + |
| 87 | + def _predict_batch(self, texts: Sequence[str]) -> list[str | None]: |
| 88 | + import torch |
| 89 | + |
| 90 | + if not self._loaded: |
| 91 | + self.load() |
| 92 | + assert self._idx2lang is not None |
| 93 | + assert self._device is not None |
| 94 | + |
| 95 | + results: list[str | None] = [] |
| 96 | + for start in range(0, len(texts), self._INTERNAL_BATCH): |
| 97 | + chunk = list(texts[start : start + self._INTERNAL_BATCH]) |
| 98 | + batch = self._encode(chunk).to(self._device) |
| 99 | + with torch.no_grad(): |
| 100 | + logits = self._model(batch) |
| 101 | + pred_idx = logits.argmax(dim=-1).cpu().tolist() |
| 102 | + results.extend(self._idx2lang[int(i)] for i in pred_idx) |
| 103 | + return results |
| 104 | + |
| 105 | + def discover_supported_languages(self) -> frozenset[str]: |
| 106 | + """Return every ISO 639-3 code in the model's ``lang2idx`` map.""" |
| 107 | + if not self._loaded: |
| 108 | + self.load() |
| 109 | + assert self._idx2lang is not None |
| 110 | + codes: set[str] = set() |
| 111 | + for code in self._idx2lang.values(): |
| 112 | + conformed = self._conform(code) |
| 113 | + if conformed is not None: |
| 114 | + codes.add(conformed) |
| 115 | + return frozenset(codes) |
0 commit comments