translator/preprocessor.py at main · findatalab/translator · GitHub

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
134
135
136
137
138
"""
Предобработка текста и выравнивание предложений
"""

import re
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
from transformers import AutoTokenizer, AutoModel
import torch

class Preprocessor:
    def __init__(self):
        print("Загрузка моделей для предобработки...")
        self.model_available = False
        self.model = None
        self.tokenizer = None

        try:
            from sentence_transformers import SentenceTransformer
            self.sbert_model = SentenceTransformer('distiluse-base-multilingual-cased-v2')
            print("✓ Загружена модель для выравнивания: distiluse-base-multilingual-cased-v2")
            self.model_available = True
        except Exception as e:
            print(f"✗ Не удалось загрузить SentenceTransformer: {e}")
            self.model_available = False

    def clean_text(self, text):
        text = re.sub(r'\s+', ' ', text)
        text = re.sub(r'[^\w\s\.\,\!\?\-\'\"]', '', text)
        return text.strip()

    def tokenize_sentences(self, text):
        sentences = re.split(r'[.!?;:]+', text)
        return [s.strip() for s in sentences if s.strip()]

    def get_embedding(self, text):
        if not self.model_available:
            return None
        try:
            return self.sbert_model.encode(text)
        except Exception as e:
            print(f"Ошибка получения эмбеддинга: {e}")
            return None

    def align_sentences_simple(self, original_sentences, translated_sentences):
        aligned = []
        min_len = min(len(original_sentences), len(translated_sentences))

        for i in range(min_len):
            aligned.append((
                original_sentences[i],
                translated_sentences[i],
                1.0
            ))

        for i in range(min_len, len(original_sentences)):
            aligned.append((original_sentences[i], None, 0.0))

        return aligned

    def align_sentences_semantic(self, original_sentences, translated_sentences):
        if not self.model_available or len(original_sentences) == 0 or len(translated_sentences) == 0:
            return self.align_sentences_simple(original_sentences, translated_sentences)

        orig_embeds = []
        trans_embeds = []

        for sent in original_sentences:
            emb = self.get_embedding(sent)
            if emb is not None:
                orig_embeds.append(emb)
            else:
                orig_embeds.append(np.random.randn(512))

        for sent in translated_sentences:
            emb = self.get_embedding(sent)
            if emb is not None:
                trans_embeds.append(emb)
            else:
                trans_embeds.append(np.random.randn(512))

        if len(orig_embeds) == 0 or len(trans_embeds) == 0:
            return self.align_sentences_simple(original_sentences, translated_sentences)

        orig_embeds = np.array(orig_embeds)
        trans_embeds = np.array(trans_embeds)

        similarity = cosine_similarity(orig_embeds, trans_embeds)

        aligned = []
        used_trans = set()

        for i, orig_sent in enumerate(original_sentences):
            best_idx = -1
            best_score = 0

            search_range = range(max(0, i-3), min(len(translated_sentences), i+4))

            for j in search_range:
                if j not in used_trans and j < len(similarity[i]):
                    score = float(similarity[i, j])
                    if score > best_score and score > 0.4:
                        best_score = score
                        best_idx = j

            if best_idx >= 0:
                aligned.append((original_sentences[i], translated_sentences[best_idx], best_score))
                used_trans.add(best_idx)
            else:
                aligned.append((original_sentences[i], None, 0.0))

        return aligned

    def preprocess_pair(self, original, translated):
        orig_clean = self.clean_text(original)
        trans_clean = self.clean_text(translated)

        orig_sentences = self.tokenize_sentences(orig_clean)
        trans_sentences = self.tokenize_sentences(trans_clean)

        aligned = self.align_sentences_semantic(orig_sentences, trans_sentences)

        translation_map = {}
        for orig, trans, score in aligned:
            if trans is not None:
                translation_map[orig.lower().strip()] = trans

        return {
            'original_clean': orig_clean,
            'translated_clean': trans_clean,
            'original_sentences': orig_sentences,
            'translated_sentences': trans_sentences,
            'aligned': aligned,
            'translation_map': translation_map
        }

    def get_model_available(self):
        return self.model_available