Add RoFormer-slav

README.md

@@ -1,3 +1,73 @@
 ---
+language:
+  - orv
+  - cu
+tags:
+  - masked-language-modeling
+  - old-slavonic
+  - old-russian
+  - birchbark
+  - historical-nlp
+  - roformer
+  - rope
+  - bpe
 license: apache-2.0
 ---
+
+# RoFormer-slav
+
+A masked language model trained from scratch on Old East Slavic and Old Church Slavonic texts,
+using a RoFormer architecture with BPE tokenisation. Based on [mini-roformer-ancient-rus-v2](https://huggingface.co/AlexSychovUN/mini-roformer-ancient-rus-v2).
+
+Note: BPE token boundaries do not always align with lacuna boundaries in editorial markup,
+which inflates span-level CER. For character-level restoration tasks consider using
+[DualEmbLM](https://huggingface.co/MaximEremeev/DualEmb-slav) instead.
+
+## Architecture
+
+- **Tokenisation**: BPE (Byte Pair Encoding), vocabulary size 50k
+- **Architecture**: RoFormer encoder with Rotary Position Embeddings (RoPE)
+- **Size**: 6 layers, hidden size 512, 8 attention heads
+
+
+## Training
+
+The model was trained on a corpus of Old Russian and Church Slavonic texts assembled from the following sources:
+
+| Source | Language | Word Tokens | Link |
+|--------|----------|--------|------|
+| Birchbark manuscripts | Old Novgorodian (mostly) | 21,464 | [gramoty.ru](https://gramoty.ru) |
+| Epigraphy | Old Church Slavonic (mostly) | 8,102 | [epigraphica.ru](https://epigraphica.ru) |
+| DIACU | Old Church Slavonic; Church Slavonic (Old Russian, Middle Bulgarian, Serbian, Resava recensions); Middle Russian | 1,683,307 | [ACL Anthology](https://aclanthology.org/2025.bsnlp-1.12/) |
+| TOROT | Old Russian; Church Slavonic | 682,430 | [torottreebank.github.io](https://torottreebank.github.io) |
+| Bible (Ponomar) | Church Slavonic | 603,047 | [GitHub](https://github.com/typiconman/ponomar/tree/master/Ponomar/languages/cu/bible/elis) |
+| Byliny | Old Russian (XI–XVII c.) | 430,103 | [rusneb.ru](https://rusneb.ru/catalog/000199_000009_003636356/) |
+| Pushkin House | Old Russian | 256,503 | [lib2.pushkinskijdom.ru](https://lib2.pushkinskijdom.ru) |
+| Military Statute (Part 2) | Old Russian | 49,787 | [rusneb.ru](https://rusneb.ru/catalog/000199_000009_004093983/) |
+| NKRYA (historical) | Old Russian; Old Rus (XI–XVIII c.) | 42,412 | [ruscorpora.ru](https://ruscorpora.ru) |
+
+Masking details: MLM probability 8%, span masking, edge masking, random gap augmentation.
+
+## Usage
+
+```python
+from transformers import AutoTokenizer, AutoModelForMaskedLM
+
+tokenizer = AutoTokenizer.from_pretrained(
+    "MaximEremeev/RoFormer-slav",
+    trust_remote_code=True,
+)
+model = AutoModelForMaskedLM.from_pretrained(
+    "MaximEremeev/RoFormer-slav",
+    trust_remote_code=True,
+)
+```
+
+## Tasks
+
+- **Generated lacunae restoration** (Test A Hit@1: 0.267, CER: 0.839)
+- **Real lacunae restoration** (Test B char Hit@1: 0.158, span Hit@1: 0.063)
+
+## Contact
+
+Maxim Eremeev, maeremeev@edu.hse.ru

aeneas_decode.py

@@ -0,0 +1,234 @@
+"""
+aeneas_decode.py
+~~~~~~~~~~~~~~~~
+Confidence-ordered beam search, following Assael et al. (2022) «Ithaca/Aeneas».
+
+Algorithm per step:
+  1. Forward pass — получаем логиты для ВСЕХ оставшихся [MASK] позиций сразу.
+  2. Находим позицию с максимальной уверенностью модели
+     (argmax по max-probability среди всех масок в этом биме).
+  3. Расширяем только эту позицию: берём top-k токенов.
+  4. Обрезаем до top-k бимов по суммарному log-probability.
+  5. Повторяем пока не останется ни одной маски.
+"""
+
+from pathlib import Path
+import math
+from dataclasses import dataclass, field
+from typing import List, Optional
+
+import torch
+from transformers import PreTrainedModel, PreTrainedTokenizerBase
+
+
+@dataclass
+class Beam:
+    input_ids: torch.Tensor          # [seq_len]
+    log_prob:  float = 0.0
+    # Список (position, token_id) в порядке заполнения
+    filled:    List[tuple] = field(default_factory=list)
+
+
+def aeneas_beam_search(
+    input_ids: torch.Tensor,          # [seq_len], уже на нужном device
+    model: PreTrainedModel,
+    tokenizer: PreTrainedTokenizerBase,
+    *,
+    beam_width: int = 5,
+    temperature: float = 1.0,
+    banned_token_ids: Optional[List[int]] = None,   # e.g. [GAP_ID]
+) -> List[Beam]:
+    """
+    Возвращает `beam_width` бимов, отсортированных по убыванию log-probability.
+    """
+    device = input_ids.device
+    mask_id = tokenizer.mask_token_id
+    banned = set(banned_token_ids or [])
+
+    # Инициализируем один начальный бим
+    beams: List[Beam] = [Beam(input_ids=input_ids.clone())]
+
+    # Считаем сколько масок нужно заполнить
+    n_masks = (input_ids == mask_id).sum().item()
+
+    with torch.no_grad():
+        for _ in range(n_masks):
+            candidates: List[Beam] = []
+
+            for beam in beams:
+                mask_positions = (beam.input_ids == mask_id).nonzero(
+                    as_tuple=True)[0].tolist()
+
+                if not mask_positions:
+                    candidates.append(beam)
+                    continue
+
+                # ── Forward pass ───────────────────────────────────────────
+                logits = model(
+                    beam.input_ids.unsqueeze(0)
+                ).logits[0]               # [seq_len, vocab]
+
+                # ── Находим самую уверенную позицию ───────────────────────
+                # Для каждой маски берём вероятность наиболее вероятного токена.
+                best_pos   = None
+                best_conf  = -1.0
+
+                for pos in mask_positions:
+                    pos_logits = logits[pos] / max(temperature, 1e-6)
+                    if banned:
+                        pos_logits = pos_logits.clone()
+                        for tid in banned:
+                            if tid < pos_logits.shape[-1]:
+                                pos_logits[tid] = float("-inf")
+                    max_prob = pos_logits.softmax(dim=-1).max().item()
+                    if max_prob > best_conf:
+                        best_conf = max_prob
+                        best_pos  = pos
+
+                # ── Расширяем именно эту позицию ──────────────────────────
+                pos_logits = logits[best_pos] / max(temperature, 1e-6)
+                if banned:
+                    pos_logits = pos_logits.clone()
+                    for tid in banned:
+                        if tid < pos_logits.shape[-1]:
+                            pos_logits[tid] = float("-inf")
+
+                probs    = pos_logits.softmax(dim=-1)
+                top_probs, top_ids = probs.topk(beam_width)
+
+                for prob, token_id in zip(top_probs.tolist(), top_ids.tolist()):
+                    if prob <= 0:
+                        continue
+                    new_ids = beam.input_ids.clone()
+                    new_ids[best_pos] = token_id
+                    candidates.append(Beam(
+                        input_ids = new_ids,
+                        log_prob  = beam.log_prob + math.log(prob + 1e-12),
+                        filled    = beam.filled + [(best_pos, token_id)],
+                    ))
+
+            # ── Pruning: оставляем top-beam_width бимов ───────────────────
+            beams = sorted(candidates, key=lambda b: b.log_prob, reverse=True)
+            beams = beams[:beam_width]
+
+    return beams
+
+
+# ── Вспомогательная ф��нкция: декодирование результатов ─────────────────────────
+
+def decode_beams(
+    beams: List[Beam],
+    original_ids: torch.Tensor,
+    tokenizer: PreTrainedTokenizerBase,
+) -> List[dict]:
+    """
+    Превращает бимы в читаемый список словарей.
+
+    Возвращает:
+        [
+          {
+            "text":          полностью восстановленный текст,
+            "filled_tokens": [(position, token_str), ...] в порядке заполнения,
+            "score":         нормализованная вероятность (0..1),
+            "log_prob":      суммарный log-prob,
+          },
+          ...
+        ]
+    """
+    results = []
+    # Нормализуем вероятности через softmax по log_prob бимов
+    log_probs = torch.tensor([b.log_prob for b in beams], dtype=torch.float)
+    scores    = log_probs.softmax(dim=0).tolist()
+
+    for beam, score in zip(beams, scores):
+        text = tokenizer.decode(beam.input_ids, skip_special_tokens=True)
+
+        filled_tokens = [
+            (pos, tokenizer.decode([tid], skip_special_tokens=True,
+                                   clean_up_tokenization_spaces=False).strip())
+            for pos, tid in beam.filled
+        ]
+
+        results.append({
+            "text":          text,
+            "filled_tokens": filled_tokens,
+            "score":         round(score, 4),
+            "log_prob":      round(beam.log_prob, 4),
+        })
+
+    return results
+
+
+# ── Высокоуровневый интерфейс ───────────────────────────────────────────────────
+
+def restore(
+    text: str,
+    model: PreTrainedModel,
+    tokenizer: PreTrainedTokenizerBase,
+    *,
+    beam_width: int = 5,
+    temperature: float = 1.0,
+    gap_token: str = "[GAP]",
+    max_length: int = 512,
+) -> List[dict]:
+    """
+    Высокоуровневая обёртка: принимает строку с [MASK], возвращает список бимов.
+
+    Args:
+        text:        текст с одним или несколькими [MASK] токенами.
+        gap_token:   токен пропуска — исключается из предсказаний.
+        beam_width:  число бимов.
+        temperature: <1 делает распределение острее, >1 — мягче.
+    """
+    device = next(model.parameters()).device
+
+    enc = tokenizer(
+        text,
+        return_tensors="pt",
+        truncation=True,
+        max_length=max_length,
+    )
+    input_ids = enc["input_ids"][0].to(device)
+
+    # Исключаем [GAP] из предсказаний
+    banned = []
+    if gap_token in tokenizer.get_vocab():
+        banned.append(tokenizer.convert_tokens_to_ids(gap_token))
+
+    beams = aeneas_beam_search(
+        input_ids, model, tokenizer,
+        beam_width=beam_width,
+        temperature=temperature,
+        banned_token_ids=banned,
+    )
+
+    return decode_beams(beams, input_ids, tokenizer)
+
+
+# ── CLI / быстрая проверка ─────────────────────────────────────────────────────
+
+if __name__ == "__main__":
+    import argparse
+    from transformers import AutoModelForMaskedLM, AutoTokenizer
+
+    parser = argparse.ArgumentParser()
+    _HERE = Path(__file__).resolve().parent
+    parser.add_argument("--model",  default=str(_HERE / "outputs/final_model"))
+    parser.add_argument("--text",   default="поклоне ѿ [MASK] к ѥва про [MASK] ѡкупи")
+    parser.add_argument("--top_k",  type=int,   default=5)
+    parser.add_argument("--temp",   type=float, default=1.0)
+    args = parser.parse_args()
+
+    device    = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    tokenizer = AutoTokenizer.from_pretrained(args.model)
+    model     = AutoModelForMaskedLM.from_pretrained(args.model).to(device)
+    model.eval()
+
+    print(f"\nВход: {args.text}\n")
+    results = restore(args.text, model, tokenizer,
+                      beam_width=args.top_k, temperature=args.temp)
+
+    for i, r in enumerate(results, 1):
+        print(f"  [{i}] score={r['score']:.3f}  log_prob={r['log_prob']:.3f}")
+        print(f"       {r['text']}")
+        print(f"       заполнено: {r['filled_tokens']}")

config.json

@@ -0,0 +1,28 @@
+{
+  "architectures": [
+    "RoFormerForMaskedLM"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "embedding_size": 512,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 512,
+  "initializer_range": 0.02,
+  "intermediate_size": 2048,
+  "layer_norm_eps": 1e-12,
+  "max_position_embeddings": 514,
+  "model_type": "roformer",
+  "num_attention_heads": 8,
+  "num_hidden_layers": 6,
+  "pad_token_id": 1,
+  "rotary_value": false,
+  "torch_dtype": "float32",
+  "transformers_version": "4.48.0",
+  "type_vocab_size": 2,
+  "use_cache": true,
+  "vocab_size": 50000,
+  "auto_map": {
+    "AutoConfig": "configuration_roformer.RoFormerConfig",
+    "AutoModelForMaskedLM": "modeling_roformer.RoFormerForMaskedLM"
+  }
+}

configuration_roformer.py

@@ -0,0 +1,18 @@
+from transformers import RoFormerConfig
+
+
+def get_roformer_config(vocab_size: int, pad_token_id: int):
+    """
+    Returns the RoFormer configuration with parameters optimized for Ancient Russian text.
+    """
+    return RoFormerConfig(
+        vocab_size=vocab_size,
+        embedding_size=512,
+        hidden_size=512,
+        num_hidden_layers=6,
+        num_attention_heads=8,
+        intermediate_size=2048,
+        max_position_embeddings=514,
+        pad_token_id=pad_token_id,
+        rotary_value=False,
+    )

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:52b96dbc3614a8a7964630f05c314e57336c043e79ebff6bc815353a0c0ac270
+size 179464480

modeling_roformer.py

@@ -0,0 +1,11 @@
+from config import get_roformer_config
+from transformers import RoFormerForMaskedLM
+
+
+def get_model(vocab_size: int, pad_token_id: int):
+    """
+    Initializes and returns a RoFormer model for Masked Language Modeling.
+    """
+    config = get_roformer_config(vocab_size, pad_token_id)
+    model = RoFormerForMaskedLM(config)
+    return model

special_tokens_map.json

@@ -0,0 +1,18 @@
+{
+  "additional_special_tokens": [
+    {
+      "content": "[GAP]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false
+    }
+  ],
+  "bos_token": "<s>",
+  "cls_token": "<s>",
+  "eos_token": "</s>",
+  "mask_token": "<mask>",
+  "pad_token": "<pad>",
+  "sep_token": "</s>",
+  "unk_token": "<unk>"
+}

tokenizer_config.json

@@ -0,0 +1,67 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<pad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "4": {
+      "content": "<mask>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "5": {
+      "content": "[GAP]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [
+    "[GAP]"
+  ],
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": false,
+  "cls_token": "<s>",
+  "eos_token": "</s>",
+  "extra_special_tokens": {},
+  "mask_token": "<mask>",
+  "max_len": 512,
+  "model_max_length": 512,
+  "pad_token": "<pad>",
+  "sep_token": "</s>",
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "unk_token": "<unk>"
+}