| import re |
| import torch |
| from transformers import AutoModelForSequenceClassification, AutoTokenizer |
|
|
|
|
| URL_PATTERN = re.compile(r"https?://\S+|www\.\S+", flags=re.IGNORECASE) |
|
|
| def fast_detect_unknown(text: str) -> bool: |
| s = text.strip() |
| if not s: |
| |
| return True |
| if URL_PATTERN.search(s): |
| |
| return True |
|
|
| total = len(s) |
| alpha = sum(ch.isalpha() for ch in s) |
| digits = sum(ch.isdigit() for ch in s) |
| spaces = sum(ch.isspace() for ch in s) |
| symbols = total - alpha - digits - spaces |
| non_space = max(1, total - spaces) |
|
|
| |
| if digits / non_space >= 0.8: |
| return True |
| |
| if symbols / non_space >= 0.45: |
| return True |
| |
| if total >= 6 and (alpha / non_space) < 0.2: |
| return True |
| |
| if " " not in s and total >= 12 and (alpha / non_space) < 0.45 and (digits > 0 or symbols > 0): |
| return True |
| return False |
|
|
| def predict_with_automodel( |
| text: str, |
| model: AutoModelForSequenceClassification, |
| tokenizer: AutoTokenizer, |
| max_length: int = 128, |
| ) -> tuple[str, float]: |
| encoded = tokenizer( |
| text, |
| return_tensors="pt", |
| truncation=True, |
| max_length=max_length, |
| ).to(model.device) |
| with torch.no_grad(): |
| logits = model(**encoded).logits |
| probs = torch.softmax(logits, dim=-1).squeeze(0) |
| pred_id = int(torch.argmax(probs).item()) |
| pred_score = float(probs[pred_id].item()) |
| pred_label = model.config.id2label[pred_id] |
| return pred_label, pred_score |
|
|
|
|
| if __name__ == "__main__": |
| model_dir = "chiennv/langid-mmbert-small" |
| tokenizer = AutoTokenizer.from_pretrained(model_dir) |
| model = AutoModelForSequenceClassification.from_pretrained(model_dir, dtype=torch.float16, device_map="auto") |
| model.eval() |
| texts = [ |
| "This is a test sentence.", |
| "Bonjour tout le monde", |
| "https://example123.com/p/9999", |
| "Christiano Ronaldo", |
| "1234567890", |
| "1234567890 is a phone number", |
| "Hello", |
| "A footballer", |
| "a dog is barking", |
| "le chatte", |
| "le chatte est un animal", |
| "고양이는 동물입니다.", |
| "猫は動物です。", |
| ] |
| for t in texts: |
| if fast_detect_unknown(t): |
| print(t, "->", "UNKNOWN", 1.0) |
| continue |
| label, score = predict_with_automodel(t, model=model, tokenizer=tokenizer) |
| print(t, "->", label, round(score, 4)) |
