Add files via upload

Author: Den4ikAI

README.md

@@ -0,0 +1,37 @@
+# RUAccent
+
+RUAccent - это библиотека для автоматической расстановки ударений на русском языке. 
+
+## Установка
+   **Требуется установленный GIT**
+   ```
+   pip install git+https://github.com/Den4ikAI/ruaccent.git
+   ```
+## Методы
+
+RUAccent предоставляет следующие методы:
+
+- `load(omograph_model_size='medium', dict_load_startup=False)`: Загрузка моделей и словарей. На данные момент доступны две модели: medium    (рекомендуется к использованию) и small. Переменная dict_load_startup отвечает за загрузку всего словаря (требуется больше ОЗУ), либо во время работы для необходимых слов (экономит ОЗУ, но требует быстрые ЖД и работает медленее)
+
+- `process_all(text)`: Обрабатывает текст всем сразу (ёфикация, расстановка ударений и расстановка ударений в словах-омографах)
+
+- `process_omographs(text)`: Расстановка ударений только в омографах.
+
+- `process_yo(text)`: Ёфикация текста.
+
+## Пример использования
+```python
+from ruaccent import RUAccent
+
+accentizer = RUAccent()
+accentizer.load(omograph_model_size='medium', dict_load_startup=False)
+
+text = 'на двери висит замок'
+print(text_processor.process_all(text))
+
+text = 'ежик нашел в лесу ягоды'
+print(text_processor.process_yo(text))
+```
+
+
+Файлы моделей и словарей располагаются по [ссылке](https://huggingface.co/TeraTTS/accentuator). Датасеты будут скоро опубликованы. Мы будем признательны, если вы будете расширять словари и загружать их в репозиторий. Это поможет улучшать данный проект.

ruaccent/__init__.py

@@ -0,0 +1 @@
+from .ruaccent import RUAccent

ruaccent/accent_model.py

@@ -0,0 +1,23 @@
+import torch
+from .char_tokenizer import CharTokenizer
+from transformers import AlbertForTokenClassification
+
+class AccentModel:
+    def __init__(self) -> None:
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    def load(self, path):
+        self.model = AlbertForTokenClassification.from_pretrained(path).to(self.device)
+        self.tokenizer = CharTokenizer.from_pretrained(path)
+    
+    def render_stress(self, word, index):
+        word = list(word)
+        word[index-1] = '+' + word[index-1]
+        return ''.join(word)
+    
+    def put_accent(self, word):
+        inputs = self.tokenizer(word, return_tensors="pt").to(self.device)
+        with torch.no_grad():
+            logits = self.model(**inputs).logits
+            predictions = torch.argmax(logits, dim=2)
+            predicted_token_class = [self.model.config.id2label[t.item()] for t in predictions[0]]
+        return self.render_stress(word, predicted_token_class.index('STRESS'))

ruaccent/char_tokenizer.py

@@ -0,0 +1,112 @@
+import os
+from typing import Optional, Tuple, List
+from collections import OrderedDict
+
+from transformers import PreTrainedTokenizer
+
+
+def load_vocab(vocab_file):
+    vocab = OrderedDict()
+    with open(vocab_file, "r", encoding="utf-8") as reader:
+        tokens = reader.readlines()
+    for index, token in enumerate(tokens):
+        token = token.rstrip("\n")
+        vocab[token] = index
+    return vocab
+
+
+class CharTokenizer(PreTrainedTokenizer):
+    vocab_files_names = {"vocab_file": "vocab.txt"}
+
+    def __init__(
+        self,
+        vocab_file=None,
+        pad_token="[pad]",
+        unk_token="[unk]",
+        bos_token="[bos]",
+        eos_token="[eos]",
+        do_lower_case=False,
+        *args,
+        **kwargs
+    ):
+        super().__init__(
+            pad_token=pad_token,
+            unk_token=unk_token,
+            bos_token=bos_token,
+            eos_token=eos_token,
+            do_lower_case=do_lower_case,
+            **kwargs
+        )
+        self.do_lower_case = do_lower_case
+
+        if not vocab_file or not os.path.isfile(vocab_file):
+            self.vocab = OrderedDict()
+            self.ids_to_tokens = OrderedDict()
+        else:
+            self.vocab = load_vocab(vocab_file)
+            self.ids_to_tokens = OrderedDict([(ids, tok) for tok, ids in self.vocab.items()])
+
+    @property
+    def vocab_size(self):
+        return len(self.vocab)
+
+    def get_vocab(self):
+        return self.vocab
+
+    def _convert_token_to_id(self, token):
+        if self.do_lower_case:
+            token = token.lower()
+        return self.vocab.get(token, self.vocab[self.unk_token])
+
+    def _convert_id_to_token(self, index):
+        return self.ids_to_tokens[index]
+
+    def _tokenize(self, text):
+        if self.do_lower_case:
+            text = text.lower()
+        return list(text)
+
+    def convert_tokens_to_string(self, tokens):
+        return "".join(tokens)
+
+    def build_inputs_with_special_tokens(
+        self,
+        token_ids_0: List[int],
+        token_ids_1: Optional[List[int]] = None
+    ) -> List[int]:
+        bos = [self.bos_token_id]
+        eos = [self.eos_token_id]
+        return bos + token_ids_0 + eos
+
+    def get_special_tokens_mask(
+         self,
+         token_ids_0: List[int],
+         token_ids_1: Optional[List[int]] = None
+    ) -> List[int]:
+        return [1] + ([0] * len(token_ids_0)) + [1]
+
+    def create_token_type_ids_from_sequences(
+        self,
+        token_ids_0: List[int],
+        token_ids_1: Optional[List[int]] = None
+    ) -> List[int]:
+        return (len(token_ids_0) + 2) * [0]
+
+    def save_vocabulary(
+        self,
+        save_directory: str,
+        filename_prefix: Optional[str] = None
+    ) -> Tuple[str]:
+        assert os.path.isdir(save_directory)
+        vocab_file = os.path.join(
+            save_directory,
+            (filename_prefix + "-" if filename_prefix else "") +
+            self.vocab_files_names["vocab_file"]
+        )
+        index = 0
+        with open(vocab_file, "w", encoding="utf-8") as writer:
+            for token, token_index in sorted(self.vocab.items(), key=lambda kv: kv[1]):
+                assert index == token_index
+                writer.write(token + "\n")
+                index += 1
+        return (vocab_file,)

ruaccent/omograph_model.py

@@ -0,0 +1,21 @@
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+import torch
+
+class OmographModel:
+    def __init__(self) -> None:
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        
+    def load(self, path):
+        self.nli_model = AutoModelForSequenceClassification.from_pretrained(path, torch_dtype=torch.bfloat16).to(self.device)
+        self.tokenizer = AutoTokenizer.from_pretrained(path)
+        
+    def classify(self, text, hypotheses):
+        encodings = self.tokenizer.batch_encode_plus([(text, hyp) for hyp in hypotheses], return_tensors='pt', padding=True)
+        input_ids = encodings['input_ids'].to(self.device)
+        with torch.no_grad():
+            logits = self.nli_model(input_ids)[0]
+            entail_contradiction_logits = logits[:,[0,2]]
+            probs = entail_contradiction_logits.softmax(dim=1)
+            prob_label_is_true = [float(p[1]) for p in probs]
+
+        return hypotheses[prob_label_is_true.index(max(prob_label_is_true))]

ruaccent/ruaccent.py

@@ -0,0 +1,100 @@
+import json
+import pathlib
+from huggingface_hub import snapshot_download
+import os
+from .omograph_model import OmographModel
+from .accent_model import AccentModel
+import re
+
+class RUAccent:
+    def __init__(self):
+        self.omograph_model = OmographModel()
+        self.accent_model = AccentModel()
+        self.workdir = str(pathlib.Path(__file__).resolve().parent)
+
+    def load(self, omograph_model_size='medium', dict_load_startup=False, repo="TeraTTS/accentuator"):
+        if not os.path.exists(self.workdir + '/dictionary') or not os.path.exists(self.workdir + '/nn'):
+            snapshot_download(repo_id=repo, ignore_patterns=["*.md", '*.gitattributes'], local_dir=self.workdir)
+        self.omographs = json.load(open(self.workdir + '/dictionary/omographs.json'))
+        self.yo_words = json.load(open(self.workdir + '/dictionary/yo_words.json'))
+        self.dict_load_startup = dict_load_startup
+        if dict_load_startup:
+            self.accents = json.load(open(self.workdir + '/dictionary/accents.json'))
+        if omograph_model_size not in ['small', 'medium']:
+            raise NotImplementedError 
+        self.omograph_model.load(self.workdir + f'/nn/nn_omograph/{omograph_model_size}/')
+        self.accent_model.load(self.workdir + '/nn/nn_accent/')
+
+    def split_by_words(self, text):
+        text = text.lower()
+        spec_chars = '!"#$%&\'()*,-./:;<=>?@[\\]^_`{|}~\r\n\xa0«»\t—…'
+        text = re.sub('[' + spec_chars + ']', ' ', text)
+        text = re.sub(' +', ' ', text)
+        output = text.split()
+        return output
+    
+    def extract_initial_letters(self, text):
+        words = self.split_by_words(text)
+        initial_letters = []
+        for word in words:
+            if len(word) > 2:
+                initial_letters.append(word[0])
+
+        return initial_letters
+    
+    def load_dict(self, text):
+        chars = self.extract_initial_letters(text)
+        out_dict = {}
+        for char in chars:
+            out_dict.update(json.load(open(f'{self.workdir}/dictionary/letter_accent/{char}.json')))
+        return out_dict
+
+    def process_punc(self, original_text, processed_text):
+        original_text = self.split_by_words(original_text)
+        processed_text = self.split_by_words(processed_text)
+        for i, word_to_process in enumerate(original_text):
+            spec_chars = 'абвгдеёжзийклмнопрстухфцчшщъыьэюя'
+            word_to_append = re.sub('[' + spec_chars + ']', ' ', word_to_process)
+            processed_text[i] = processed_text[i] + word_to_append.strip()
+        return ' '.join(processed_text)
+
+    def count_vowels(self, text):
+        vowels = 'аеёиоуыэюяАЕЁИОУЫЭЮЯ'
+        return sum(1 for char in text if char in vowels)
+
+    def process_omographs(self, text):
+        splitted_text = self.split_by_words(text)
+        founded_omographs = []
+        for i, word in enumerate(splitted_text):
+            variants = self.omographs.get(word)
+            if variants:
+                founded_omographs.append({'word': word, 'variants': variants, 'position': i})
+        for omograph in founded_omographs:
+            splitted_text[omograph['position']] = f"<w>{splitted_text[omograph['position']]}</w>"
+            cls = self.omograph_model.classify(' '.join(splitted_text), omograph['variants'])
+            splitted_text[omograph['position']] = cls
+        return ' '.join(splitted_text)
+
+    def process_yo(self, text):
+        splitted_text = self.split_by_words(text)
+        for i, word in enumerate(splitted_text):
+            splitted_text[i] = self.yo_words.get(word, word)
+        return ' '.join(splitted_text)
+
+    def process_accent(self, text):
+        if not self.dict_load_startup:
+            self.accents = self.load_dict(text)
+        splitted_text = self.split_by_words(text)
+        for i, word in enumerate(splitted_text):
+            stressed_word = self.accents.get(word, word)
+            if '+' not in stressed_word and self.count_vowels(word) > 1:
+                splitted_text[i] = self.accent_model.put_accent(word)
+            else:
+                splitted_text[i] = stressed_word
+        return ' '.join(splitted_text)
+
+    def process_all(self, text):
+        processed_text = self.process_yo(text)
+        processed_text = self.process_omographs(processed_text)
+        processed_text = self.process_accent(processed_text)
+        return processed_text

setup.py

@@ -0,0 +1,27 @@
+from setuptools import setup, find_packages
+
+setup(
+    name='ruaccent',
+    version='1.0.0',
+    author='Denis Petrov',
+    author_email='arduino4b@gmail.com',
+    description='A Russian text accentuation tool',
+    license='MIT', 
+    url='https://github.com/Den4ikAI/ruaccent',  
+    packages=find_packages(),
+    install_requires=[
+        'huggingface_hub',
+        'torch==1.13.1',
+        'transformers',
+        'sentencepiece'
+    ],
+    classifiers=[
+        'Development Status :: 5 - Production/Stable',
+        'Intended Audience :: Developers',
+        'License :: OSI Approved :: MIT License',
+        'Programming Language :: Python :: 3',
+        'Operating System :: Microsoft :: Windows',
+        'Operating System :: Unix',
+        'Operating System :: MacOS',
+    ],
+)