"""MixLID inference using standard Hugging Face Transformers models.""" from collections import Counter from pathlib import Path import torch from transformers import AutoModelForTokenClassification, AutoTokenizer AUXILIARY_LABELS = {"NUM", "PUNCT", "SYM"} def _format_prediction(text, tokens): votes = Counter(t["label"] for t in tokens if t["label"] not in AUXILIARY_LABELS) language = min(votes, key=lambda label: (-votes[label], label)) if votes else "und" spans = [] counts = [] for token in tokens: if (spans and spans[-1]["label"] == token["label"] and token["start"] >= spans[-1]["end"] and not text[spans[-1]["end"]:token["start"]].strip(" ")): span = spans[-1] span["score"] = (span["score"] * counts[-1] + token["score"]) / (counts[-1] + 1) counts[-1] += 1 span["end"] = token["end"] span["text"] = text[span["start"]:span["end"]] else: spans.append(dict(token)) counts.append(1) return { "text": text, "language": language, "language_score": votes[language] / sum(votes.values()) if votes else 0.0, "tokens": tokens, "spans": spans, } class MixLID: """Token labels, merged spans, and an input-level majority language. Offsets are Python character indices into the original, unnormalized text. The default device is CPU; select a CUDA device explicitly to use a GPU. """ def __init__(self, model, tokenizer): if not tokenizer.is_fast: raise ValueError("MixLID requires a fast tokenizer for character offsets") self.model = model.eval() self.tokenizer = tokenizer @classmethod def from_pretrained(cls, model_name_or_path=None, *, device="cpu", dtype=None, attn_implementation="sdpa", revision=None, local_files_only=False): """Load a local directory or a Hub model ID without remote model code. dtype defaults to float32 on CPU and float16 on CUDA. Pass torch.bfloat16 explicitly when appropriate for your GPU. """ path = str(model_name_or_path or Path(__file__).resolve().parent) device = torch.device(device) dtype = dtype or (torch.float16 if device.type == "cuda" else torch.float32) common = dict(revision=revision, local_files_only=local_files_only, trust_remote_code=False) tokenizer = AutoTokenizer.from_pretrained(path, use_fast=True, **common) model = AutoModelForTokenClassification.from_pretrained( path, dtype=dtype, attn_implementation=attn_implementation, reference_compile=False, **common, ).to(device) return cls(model, tokenizer) def predict_sentence(self, texts): """One majority language per input; no automatic sentence splitting. Returns a dict for a string, or a list of dicts for a list of strings. score is the language vote share, not a calibrated probability. """ results = self.predict(texts, batch_size=1) def sentence(result): return {"text": result["text"], "language": result["language"], "score": result["language_score"]} return sentence(results) if isinstance(texts, str) else [sentence(r) for r in results] def predict_tokens(self, texts): """Return subword predictions and merged language spans for each input.""" results = self.predict(texts, batch_size=1) def tokens(result): return {key: result[key] for key in ("text", "tokens", "spans")} return tokens(results) if isinstance(texts, str) else [tokens(r) for r in results] @torch.inference_mode() def predict(self, texts, *, batch_size=8, max_length=8192, stride=64): """Predict one string or a list of strings, without dropping long-input tails. batch_size counts windows. Special/padding tokens and spans containing only ordinary spaces (U+0020) are omitted. Newlines/tabs are retained. Overlapping windows keep the first prediction for each character range. """ single = isinstance(texts, str) if not single and not isinstance(texts, (list, tuple)): raise TypeError("texts must be a string or a list/tuple of strings") texts = [texts] if single else list(texts) if any(not isinstance(text, str) for text in texts): raise TypeError("every input must be a string") limit = min(self.tokenizer.model_max_length, self.model.config.max_position_embeddings) special_count = self.tokenizer.num_special_tokens_to_add(pair=False) if not isinstance(batch_size, int) or batch_size < 1: raise ValueError("batch_size must be a positive integer") if not isinstance(max_length, int) or not special_count < max_length <= limit: raise ValueError(f"max_length must be between {special_count + 1} and {limit}") if not isinstance(stride, int) or not 0 <= stride < max_length - special_count: raise ValueError("stride must be nonnegative and smaller than the window content length") if not texts: return [] encoded = self.tokenizer( texts, truncation=True, max_length=max_length, stride=stride, return_overflowing_tokens=True, return_offsets_mapping=True, return_special_tokens_mask=True, padding=False, ) offsets = encoded.pop("offset_mapping") mapping = encoded.pop("overflow_to_sample_mapping") special = encoded.pop("special_tokens_mask") results = [[] for _ in texts] seen = [set() for _ in texts] for begin in range(0, len(mapping), batch_size): end = min(begin + batch_size, len(mapping)) features = [{key: values[i] for key, values in encoded.items()} for i in range(begin, end)] inputs = self.tokenizer.pad(features, padding=True, return_tensors="pt").to(self.model.device) logits = self.model(**inputs).logits scores, labels = logits.float().softmax(dim=-1).max(dim=-1) scores, labels = scores.cpu().tolist(), labels.cpu().tolist() for row, window in enumerate(range(begin, end)): sample = mapping[window] text = texts[sample] for j, ((start, stop), is_special) in enumerate(zip(offsets[window], special[window])): if is_special or stop <= start or (start, stop) in seen[sample]: continue seen[sample].add((start, stop)) if not text[start:stop].strip(" "): continue results[sample].append({ "text": text[start:stop], "start": start, "end": stop, "label": self.model.config.id2label[labels[row][j]], "score": scores[row][j], }) predictions = [_format_prediction(text, tokens) for text, tokens in zip(texts, results)] return predictions[0] if single else predictions