""" preprocess.py - Скрипт для предварительной обработки датасета книг и создания векторных представлений для поисковой системы """ import os import pandas as pd import numpy as np import torch from transformers import AutoTokenizer, AutoModel import faiss import re import nltk from nltk.corpus import stopwords from nltk.tokenize import word_tokenize import argparse from tqdm import tqdm import nltk nltk.download('punkt') nltk.download('stopwords') nltk.download('punkt_tab') # Загрузим стоп-слова для русского языка try: nltk.data.find('corpora/stopwords') except LookupError: nltk.download('stopwords') try: nltk.data.find('tokenizers/punkt') except LookupError: nltk.download('punkt') stop_words = set(stopwords.words('russian')) # Класс для получения эмбеддингов с помощью RuBERT class RuBERTEmbedder: def __init__(self, model_name="DeepPavlov/rubert-base-cased"): print(f"Загрузка модели {model_name}...") self.tokenizer = AutoTokenizer.from_pretrained(model_name) self.model = AutoModel.from_pretrained(model_name) self.model.eval() # Используем GPU если доступен, иначе CPU self.device = "cuda" if torch.cuda.is_available() else "cpu" print(f"Используется устройство: {self.device}") self.model.to(self.device) def mean_pooling(self, model_output, attention_mask): """Среднее значение по токенам для получения эмбеддинга предложения""" token_embeddings = model_output[0] input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float() return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9) def get_embedding(self, text): """Получение векторного представления текста""" encoded_input = self.tokenizer(text, padding=True, truncation=True, max_length=512, return_tensors='pt') encoded_input = {k: v.to(self.device) for k, v in encoded_input.items()} with torch.no_grad(): model_output = self.model(**encoded_input) embeddings = self.mean_pooling(model_output, encoded_input['attention_mask']) return embeddings.cpu().numpy()[0] def get_embeddings_batch(self, texts, batch_size=8): """Получение векторных представлений для списка текстов с использованием батчей""" all_embeddings = [] for i in tqdm(range(0, len(texts), batch_size), desc="Создание эмбеддингов"): batch_texts = texts[i:i+batch_size] # Обработка пустых строк batch_texts = [text if text and isinstance(text, str) else " " for text in batch_texts] encoded_input = self.tokenizer(batch_texts, padding=True, truncation=True, max_length=512, return_tensors='pt') encoded_input = {k: v.to(self.device) for k, v in encoded_input.items()} with torch.no_grad(): model_output = self.model(**encoded_input) embeddings = self.mean_pooling(model_output, encoded_input['attention_mask']) all_embeddings.append(embeddings.cpu().numpy()) return np.vstack(all_embeddings) def preprocess_text(text): """Предобработка текста: удаление специальных символов, приведение к нижнему регистру, удаление стоп-слов""" if isinstance(text, str): # Приведение к нижнему регистру text = text.lower() # Удаление специальных символов text = re.sub(r'[^\w\s]', '', text) # Токенизация tokens = word_tokenize(text, language='russian') # Удаление стоп-слов filtered_tokens = [word for word in tokens if word not in stop_words] # Объединение обратно в строку return ' '.join(filtered_tokens) return '' def prepare_data(input_file, output_dir="model", annotation_column="annotation", title_column="title", author_column="author", image_url_column="image_url", page_url_column="page_url", sample_size=None): """Подготовка данных для поисковой системы""" # Создание выходной директории, если она не существует os.makedirs(output_dir, exist_ok=True) print(f"Загрузка данных из {input_file}...") df = pd.read_csv(input_file) # Проверка наличия обязательной колонки с аннотацией if annotation_column not in df.columns: raise ValueError(f"В файле отсутствует колонка с аннотациями: {annotation_column}") # Очистка от записей без аннотации initial_size = len(df) df = df.dropna(subset=[annotation_column]) print(f"Удалено записей без аннотаций: {initial_size - len(df)}") # Если указан размер выборки, отбираем случайные записи if sample_size and sample_size < len(df): df = df.sample(sample_size, random_state=42) print(f"Используется случайная выборка из {sample_size} записей") # Предобработка аннотаций print("Предобработка аннотаций...") df['processed_annotation'] = df[annotation_column].apply(preprocess_text) # Загрузка модели для векторизации print("Инициализация модели для векторизации...") embedder = RuBERTEmbedder() # Векторизация аннотаций print("Векторизация аннотаций...") annotations = df['processed_annotation'].tolist() embeddings = embedder.get_embeddings_batch(annotations) # Создание индекса FAISS print("Создание индекса FAISS...") dimension = embeddings.shape[1] index = faiss.IndexFlatIP(dimension) faiss.normalize_L2(embeddings) index.add(embeddings) # Сохранение обработанных данных и индекса print(f"Сохранение данных в {output_dir}...") # Сохраняем только нужные колонки columns_to_save = [col for col in [annotation_column, title_column, author_column, image_url_column, page_url_column, 'processed_annotation'] if col in df.columns] df[columns_to_save].to_csv(f"{output_dir}/book_data.csv", index=False) # Сохраняем эмбеддинги np.save(f"{output_dir}/embeddings.npy", embeddings) # Сохраняем индекс FAISS faiss.write_index(index, f"{output_dir}/faiss_index.bin") print(f"Данные успешно обработаны и сохранены в {output_dir}") print(f"Всего книг: {len(df)}") return df def main(): parser = argparse.ArgumentParser(description='Предобработка датасета книг для поисковой системы') parser.add_argument('--input', type=str, required=True, help='Путь к CSV файлу с данными книг') parser.add_argument('--output', type=str, default='model', help='Директория для сохранения модели и данных') parser.add_argument('--annotation', type=str, default='annotation', help='Имя колонки с аннотациями') parser.add_argument('--title', type=str, default='title', help='Имя колонки с названиями книг') parser.add_argument('--author', type=str, default='author', help='Имя колонки с авторами') parser.add_argument('--image_url', type=str, default='image_url', help='Имя колонки с URL изображений') parser.add_argument('--page_url', type=str, default='page_url', help='Имя колонки с URL страниц') parser.add_argument('--sample', type=int, default=None, help='Размер выборки (если нужно ограничить)') args = parser.parse_args() prepare_data( input_file=args.input, output_dir=args.output, annotation_column=args.annotation, title_column=args.title, author_column=args.author, image_url_column=args.image_url, page_url_column=args.page_url, sample_size=args.sample ) if __name__ == "__main__": main()