| |
| """Word-difficulty regressor for the game «Шляпа» (thehat.tech). |
| |
| Predicts the Elo-style difficulty rating E of a Russian noun from the word |
| itself plus its corpus frequency. See README.md for metrics and background. |
| |
| Usage: |
| from modeling import WordDifficultyPredictor |
| predictor = WordDifficultyPredictor.from_dir(".") |
| predictor.predict(["кровать", "соразмерность"]) # -> [~38.5, ~60.0] |
| """ |
|
|
| import json |
| import os |
|
|
| import numpy as np |
| import torch |
| from transformers import AutoModel, AutoTokenizer |
|
|
| MAX_LEN = 16 |
|
|
|
|
| class BertFreqRegressor(torch.nn.Module): |
| def __init__(self, base_model="DeepPavlov/rubert-base-cased"): |
| super().__init__() |
| self.bert = AutoModel.from_pretrained(base_model) |
| hidden = self.bert.config.hidden_size |
| self.head = torch.nn.Sequential( |
| torch.nn.Linear(hidden + 2, 128), |
| torch.nn.GELU(), |
| torch.nn.Dropout(0.1), |
| torch.nn.Linear(128, 1), |
| ) |
|
|
| def forward(self, ids, mask, extra): |
| out = self.bert(input_ids=ids, attention_mask=mask) |
| m = mask.unsqueeze(-1).float() |
| pooled = (out.last_hidden_state * m).sum(1) / m.sum(1).clamp(min=1) |
| return self.head(torch.cat([pooled, extra], dim=-1)).squeeze(-1) |
|
|
|
|
| class WordDifficultyPredictor: |
| def __init__(self, model, tokenizer, freq, mean, std, lf_mean, lf_std, |
| device=None): |
| self.model = model |
| self.tokenizer = tokenizer |
| self.freq = freq |
| self.mean, self.std = mean, std |
| self.lf_mean, self.lf_std = lf_mean, lf_std |
| self.device = device or torch.device( |
| "mps" if torch.backends.mps.is_available() |
| else "cuda" if torch.cuda.is_available() else "cpu") |
| self.model.to(self.device).eval() |
|
|
| @classmethod |
| def from_dir(cls, path): |
| """Load from a local directory (e.g. a huggingface_hub snapshot).""" |
| ckpt = torch.load(os.path.join(path, "model.pt"), map_location="cpu", |
| weights_only=True) |
| model = BertFreqRegressor(ckpt.get("model_name", |
| "DeepPavlov/rubert-base-cased")) |
| model.load_state_dict(ckpt["state_dict"]) |
| tokenizer = AutoTokenizer.from_pretrained(path) |
| with open(os.path.join(path, "word_frequency.json"), |
| encoding="utf-8") as fh: |
| freq = json.load(fh) |
| return cls(model, tokenizer, freq, ckpt["mean"], ckpt["std"], |
| ckpt["lf_mean"], ckpt["lf_std"]) |
|
|
| @classmethod |
| def from_hub(cls, repo_id="nzinov/thehat-word-difficulty"): |
| from huggingface_hub import snapshot_download |
| return cls.from_dir(snapshot_download(repo_id)) |
|
|
| @torch.no_grad() |
| def predict(self, words): |
| """Difficulty ratings E for a list of Russian words (lowercase).""" |
| words = [w.strip().lower() for w in words] |
| enc = self.tokenizer(words, truncation=True, max_length=MAX_LEN, |
| padding="max_length", return_tensors="pt") |
| feats = [] |
| for w in words: |
| f = self.freq.get(w) |
| has = 1.0 if f is not None else 0.0 |
| lf = ((np.log1p(f) - self.lf_mean) / self.lf_std |
| if f is not None else 0.0) |
| feats.append([lf, has]) |
| extra = torch.tensor(feats, dtype=torch.float32) |
| pred = self.model(enc["input_ids"].to(self.device), |
| enc["attention_mask"].to(self.device), |
| extra.to(self.device)) |
| return (pred.cpu().numpy() * self.std + self.mean).tolist() |
|
|
|
|
| if __name__ == "__main__": |
| import sys |
| predictor = WordDifficultyPredictor.from_dir( |
| os.path.dirname(os.path.abspath(__file__)) or ".") |
| words = sys.argv[1:] or ["кровать", "луна", "синоним", "соразмерность"] |
| for word, e in zip(words, predictor.predict(words)): |
| print(f"{word:20s} E = {e:.1f}") |
|
|
