📦 ibakaidov / text-corrector

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
import random
import string
import re

russian_alphabet = ''.join([chr(code) for code in range(ord('а'), ord('я') + 1)]) + \
                   ''.join([chr(code) for code in range(ord('А'), ord('Я') + 1)]) + 'ёЁ'
alphabet = string.ascii_letters + russian_alphabet + string.digits + string.punctuation + ' '
_punctuation_chars = string.punctuation + '«»„“”‚‘’‹›—–…№•·”“'
_punctuation_table = str.maketrans({ch: ' ' for ch in _punctuation_chars})

def text_prepare(text, variations=3):
    """
    Подготавливает зашумленный текст для обучения модели.
    
    :param text: str: Входной текст
    :param variations: int: Количество вариаций зашумленного текста
    """

    text = remove_punctuation(text)
    variations_list = []
    for _ in range(variations):
        noisy_text = add_noise_to_text(text, percentage=0.2)
        cleaned_text = remove_random_characters(noisy_text, percentage=0.1)
        variations_list.append(cleaned_text)    
    return variations_list

def last_word_prepare(text, variations=3):
    """
    Подготавливает текст, внося шум только в последнее слово.
    :param text: str: Входной текст 
    :param variations: int: Количество вариаций зашумленного текста
    :return: list of dicts:
        {
            "context": str,  # Текст до текущего слова
            "position": int, # Позиция вставки (0 - текущая позиция, 1 - новое слово, 2 - внутри слова)
            "prediction": str # Слово или его часть с шумом
        }

    """ 
    text = remove_punctuation(text)
    words = text.split(' ')
    return_obj = []
    if len(words) == 0:
        return return_obj
    for i in range(len(words)): 
        current_word = words[i]
        previous_words = words[:i] if i > 0 else []
        return_obj.append({
            "context": ' '.join(previous_words),
            "position": 1,
            "prediction": current_word,
        })
        return_obj.append({
            "context": ' '.join(previous_words) + ' ',
            "position": 0,
            "prediction": current_word,
        })
        for j in range(len(current_word)):
            start = ' '.join(previous_words) + ' ' if i > 0 else '' 
            start += current_word[:j]
            return_obj.append({
                "context": start,
                "position": 2,
                "prediction": current_word,
            })
            
            for k in range(variations): 
                start = ' '.join(previous_words) + ' ' if i > 0 else ''
                end = current_word[j+1:]
                noisy_char = add_noise_to_text(text=current_word[j], percentage=0.5)
                start_noisy = start + current_word[:j] + noisy_char
                return_obj.append({
                    "context": start_noisy,
                    "position": 2,
                    "prediction": current_word,
                })
            
            
    return return_obj

def add_noise_to_text(text, percentage=0.2):
    """
    Добавляет шум в текст, изменяя случайные символы.
    
    :param text: str: Входной текст
    :param percentage: float: Процент символов для зашумления
    :return: str: Текст с добавленным шумом
    """
    noisy_text = list(text)
    num_noises = max(1, int(len(text) * percentage))  # Добавляем шум в percentage символов
    
    for _ in range(num_noises):
        index = random.randint(0, len(text) - 1)
        noisy_text[index] = random.choice(alphabet)
    
    return ''.join(noisy_text)

def remove_random_characters(text, percentage=0.1):
    """
    Удаляет случайные символы из текста.
    
    :param text: str: Входной текст
    :param percentage: float: Процент символов для удаления
    :return: str: Текст с удаленными символами
    """
    text_list = list(text)
    num_removals = max(1, int(len(text) * percentage))  # Удаляем percentage символов
    
    for _ in range(num_removals):
        if len(text_list) > 1:
            index = random.randint(0, len(text_list) - 1)
            del text_list[index]
    
    return ''.join(text_list)

def remove_punctuation(text: str) -> str:
    """
    Удаляет знаки препинания и заменяет их пробелами, схлопывая множественные пробелы.
    """
    text = text.translate(_punctuation_table)
    return re.sub(r'\s+', ' ', text).strip()

if __name__ == "__main__":
    sample_text = "Пример текста для тестирования функции подготовки текста."
    variations = text_prepare(sample_text, variations=5)
    for i, var in enumerate(variations):
        print(f"Variation {i+1}: {var}")    
        
    last_word_variations = last_word_prepare(sample_text, variations=2)
    for i, var in enumerate(last_word_variations):
        print(f"Last Word Variation {i+1}: {var}")