# src/train_subword.py """ Модуль для обучения подсловных моделей токенизации (BPE, WordPiece, Unigram). Поддерживает обучение моделей с различными параметрами и их сравнительный анализ. """ import os import json import time from typing import List, Dict, Tuple, Optional, Any from dataclasses import dataclass from pathlib import Path import pandas as pd # Импорты для различных библиотек токенизации try: from tokenizers import Tokenizer, trainers, models, pre_tokenizers, normalizers from tokenizers.trainers import BpeTrainer, WordPieceTrainer, UnigramTrainer TOKENIZERS_AVAILABLE = True except ImportError: TOKENIZERS_AVAILABLE = False try: import sentencepiece as spm SENTENCEPIECE_AVAILABLE = True except ImportError: SENTENCEPIECE_AVAILABLE = False @dataclass class SubwordModelConfig: """Конфигурация для обучения подсловной модели.""" model_type: str # 'bpe', 'wordpiece', 'unigram' vocab_size: int min_frequency: int = 2 special_tokens: List[str] = None model_name: str = "" def __post_init__(self): if self.special_tokens is None: self.special_tokens = ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"] if not self.model_name: self.model_name = f"{self.model_type}_{self.vocab_size}" @dataclass class SubwordMetrics: """Метрики для оценки подсловных моделей.""" model_name: str vocab_size: int fragmentation_rate: float compression_ratio: float reconstruction_accuracy: float training_time: float oov_rate: float = 0.0 class SubwordModelTrainer: """Класс для обучения и сравнения подсловных моделей токенизации.""" def __init__(self, output_dir: str = "models"): """ Инициализация тренера. Args: output_dir: Директория для сохранения моделей """ self.output_dir = Path(output_dir) self.output_dir.mkdir(exist_ok=True) self.models = {} self.metrics = {} def prepare_corpus(self, input_path: str, output_path: str, text_field: str = 'text') -> int: """ Подготавливает корпус для обучения подсловных моделей. Args: input_path: Путь к JSONL файлу с корпусом output_path: Путь для сохранения подготовленного корпуса text_field: Поле с текстом статьи Returns: Количество обработанных статей """ import json texts = [] with open(input_path, 'r', encoding='utf-8') as f: for line in f: try: article = json.loads(line.strip()) if text_field in article and article[text_field].strip(): texts.append(article[text_field]) except json.JSONDecodeError: continue # Сохраняем корпус как текстовый файл with open(output_path, 'w', encoding='utf-8') as f: for text in texts: f.write(text + '\n') return len(texts) def train_bpe_model(self, config: SubwordModelConfig, corpus_path: str) -> str: """ Обучает BPE модель. Args: config: Конфигурация модели corpus_path: Путь к корпусу Returns: Путь к сохраненной модели """ if not TOKENIZERS_AVAILABLE: raise ImportError("Библиотека tokenizers не установлена") # Создаем токенизатор tokenizer = Tokenizer(models.BPE()) tokenizer.pre_tokenizer = pre_tokenizers.Whitespace() # Создаем тренер trainer = BpeTrainer( vocab_size=config.vocab_size, min_frequency=config.min_frequency, special_tokens=config.special_tokens ) # Обучаем модель start_time = time.time() tokenizer.train([corpus_path], trainer) training_time = time.time() - start_time # Сохраняем модель model_path = self.output_dir / f"{config.model_name}.json" tokenizer.save(str(model_path)) # Сохраняем метрики self.metrics[config.model_name] = { 'training_time': training_time, 'model_type': 'bpe' } return str(model_path) def train_wordpiece_model(self, config: SubwordModelConfig, corpus_path: str) -> str: """ Обучает WordPiece модель. Args: config: Конфигурация модели corpus_path: Путь к корпусу Returns: Путь к сохраненной модели """ if not TOKENIZERS_AVAILABLE: raise ImportError("Библиотека tokenizers не установлена") # Создаем токенизатор tokenizer = Tokenizer(models.WordPiece()) tokenizer.pre_tokenizer = pre_tokenizers.Whitespace() # Создаем тренер trainer = WordPieceTrainer( vocab_size=config.vocab_size, min_frequency=config.min_frequency, special_tokens=config.special_tokens ) # Обучаем модель start_time = time.time() tokenizer.train([corpus_path], trainer) training_time = time.time() - start_time # Сохраняем модель model_path = self.output_dir / f"{config.model_name}.json" tokenizer.save(str(model_path)) # Сохраняем метрики self.metrics[config.model_name] = { 'training_time': training_time, 'model_type': 'wordpiece' } return str(model_path) def train_unigram_model(self, config: SubwordModelConfig, corpus_path: str) -> str: """ Обучает Unigram модель. Args: config: Конфигурация модели corpus_path: Путь к корпусу Returns: Путь к сохраненной модели """ if not TOKENIZERS_AVAILABLE: raise ImportError("Библиотека tokenizers не установлена") # Создаем токенизатор tokenizer = Tokenizer(models.Unigram()) tokenizer.pre_tokenizer = pre_tokenizers.Whitespace() # Создаем тренер trainer = UnigramTrainer( vocab_size=config.vocab_size, min_frequency=config.min_frequency, special_tokens=config.special_tokens ) # Обучаем модель start_time = time.time() tokenizer.train([corpus_path], trainer) training_time = time.time() - start_time # Сохраняем модель model_path = self.output_dir / f"{config.model_name}.json" tokenizer.save(str(model_path)) # Сохраняем метрики self.metrics[config.model_name] = { 'training_time': training_time, 'model_type': 'unigram' } return str(model_path) def train_sentencepiece_model(self, config: SubwordModelConfig, corpus_path: str) -> str: """ Обучает SentencePiece модель. Args: config: Конфигурация модели corpus_path: Путь к корпусу Returns: Путь к сохраненной модели """ if not SENTENCEPIECE_AVAILABLE: raise ImportError("Библиотека sentencepiece не установлена") # Параметры для SentencePiece model_prefix = str(self.output_dir / config.model_name) # Определяем тип модели model_type_map = { 'bpe': 'bpe', 'wordpiece': 'word', # SentencePiece не поддерживает WordPiece напрямую 'unigram': 'unigram' } spm_model_type = model_type_map.get(config.model_type, 'bpe') # Параметры обучения train_args = [ f'--input={corpus_path}', f'--model_prefix={model_prefix}', f'--vocab_size={config.vocab_size}', f'--model_type={spm_model_type}', f'--character_coverage=0.9995', f'--normalization_rule_name=nfkc', f'--user_defined_symbols={",".join(config.special_tokens)}' ] # Обучаем модель start_time = time.time() spm.SentencePieceTrainer.train(' '.join(train_args)) training_time = time.time() - start_time # Сохраняем метрики self.metrics[config.model_name] = { 'training_time': training_time, 'model_type': f'sentencepiece_{spm_model_type}' } return f"{model_prefix}.model" def train_model(self, config: SubwordModelConfig, corpus_path: str, use_sentencepiece: bool = False) -> str: """ Обучает модель указанного типа. Args: config: Конфигурация модели corpus_path: Путь к корпусу use_sentencepiece: Использовать SentencePiece вместо tokenizers Returns: Путь к сохраненной модели """ print(f"Обучаем модель {config.model_name} ({config.model_type})...") if use_sentencepiece and SENTENCEPIECE_AVAILABLE: return self.train_sentencepiece_model(config, corpus_path) if config.model_type == 'bpe': return self.train_bpe_model(config, corpus_path) elif config.model_type == 'wordpiece': return self.train_wordpiece_model(config, corpus_path) elif config.model_type == 'unigram': return self.train_unigram_model(config, corpus_path) else: raise ValueError(f"Неподдерживаемый тип модели: {config.model_type}") def evaluate_model(self, model_path: str, test_texts: List[str]) -> SubwordMetrics: """ Оценивает качество обученной модели. Args: model_path: Путь к модели test_texts: Тестовые тексты Returns: Метрики модели """ if not TOKENIZERS_AVAILABLE: raise ImportError("Библиотека tokenizers не установлена") # Загружаем модель tokenizer = Tokenizer.from_file(model_path) total_tokens = 0 total_words = 0 fragmented_words = 0 reconstruction_errors = 0 for text in test_texts: # Токенизируем encoded = tokenizer.encode(text) tokens = encoded.tokens # Декодируем обратно reconstructed = tokenizer.decode(encoded.ids) # Подсчитываем метрики words = text.split() total_words += len(words) total_tokens += len(tokens) # Подсчитываем фрагментированные слова for word in words: word_tokens = tokenizer.encode(word).tokens if len(word_tokens) > 1: fragmented_words += 1 # Проверяем точность реконструкции if reconstructed.strip() != text.strip(): reconstruction_errors += 1 # Вычисляем метрики fragmentation_rate = fragmented_words / total_words if total_words > 0 else 0 compression_ratio = total_words / total_tokens if total_tokens > 0 else 1 reconstruction_accuracy = 1 - (reconstruction_errors / len(test_texts)) if test_texts else 1 model_name = Path(model_path).stem return SubwordMetrics( model_name=model_name, vocab_size=tokenizer.get_vocab_size(), fragmentation_rate=fragmentation_rate, compression_ratio=compression_ratio, reconstruction_accuracy=reconstruction_accuracy, training_time=self.metrics.get(model_name, {}).get('training_time', 0), oov_rate=0.0 # Будет вычислено отдельно ) def train_multiple_models(self, corpus_path: str, vocab_sizes: List[int] = None) -> Dict[str, str]: """ Обучает несколько моделей с разными параметрами. Args: corpus_path: Путь к корпусу vocab_sizes: Список размеров словаря Returns: Словарь {имя_модели: путь_к_модели} """ if vocab_sizes is None: vocab_sizes = [8000, 16000, 32000] model_types = ['bpe', 'wordpiece', 'unigram'] trained_models = {} for model_type in model_types: for vocab_size in vocab_sizes: config = SubwordModelConfig( model_type=model_type, vocab_size=vocab_size, min_frequency=2 ) try: model_path = self.train_model(config, corpus_path) trained_models[config.model_name] = model_path print(f"Модель {config.model_name} обучена успешно") except Exception as e: print(f"Ошибка при обучении модели {config.model_name}: {e}") return trained_models def compare_models(self, model_paths: Dict[str, str], test_texts: List[str]) -> pd.DataFrame: """ Сравнивает несколько обученных моделей. Args: model_paths: Словарь {имя_модели: путь_к_модели} test_texts: Тестовые тексты Returns: DataFrame с результатами сравнения """ results = [] for model_name, model_path in model_paths.items(): try: metrics = self.evaluate_model(model_path, test_texts) results.append({ 'Модель': model_name, 'Тип': metrics.model_name.split('_')[0], 'Размер словаря': metrics.vocab_size, 'Процент фрагментации': round(metrics.fragmentation_rate * 100, 2), 'Коэффициент сжатия': round(metrics.compression_ratio, 3), 'Точность реконструкции': round(metrics.reconstruction_accuracy * 100, 2), 'Время обучения (сек)': round(metrics.training_time, 2) }) except Exception as e: print(f"Ошибка при оценке модели {model_name}: {e}") return pd.DataFrame(results) def save_comparison_results(self, results_df: pd.DataFrame, output_path: str): """Сохраняет результаты сравнения в CSV файл.""" results_df.to_csv(output_path, index=False, encoding='utf-8') print(f"Результаты сравнения сохранены в {output_path}") def main(): """Основная функция для обучения и сравнения подсловных моделей.""" trainer = SubwordModelTrainer() # Подготавливаем корпус corpus_path = "data/corpus.txt" if not os.path.exists(corpus_path): print("Подготавливаем корпус...") articles_count = trainer.prepare_corpus("data/raw_corpus.jsonl", corpus_path) print(f"Подготовлено {articles_count} статей") # Обучаем модели print("Обучаем подсловные модели...") trained_models = trainer.train_multiple_models(corpus_path) # Загружаем тестовые тексты test_texts = [] with open(corpus_path, 'r', encoding='utf-8') as f: for i, line in enumerate(f): if i >= 100: # Берем первые 100 строк для тестирования break test_texts.append(line.strip()) # Сравниваем модели print("Сравниваем модели...") comparison_results = trainer.compare_models(trained_models, test_texts) print("\nРезультаты сравнения:") print(comparison_results) # Сохраняем результаты trainer.save_comparison_results(comparison_results, "results/subword_comparison.csv") if __name__ == "__main__": main()