# app_modified.py - Оптимизированная версия Streamlit-приложения import streamlit as st 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 os # Загрузим стоп-слова для русского языка 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"): self.tokenizer = AutoTokenizer.from_pretrained(model_name) self.model = AutoModel.from_pretrained(model_name) self.model.eval() # Используем CPU для запуска в Spaces self.device = "cpu" 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 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 '' # Класс поисковой системы class BookSearchEngine: def __init__(self, embedder=None): self.embedder = embedder self.faiss_index = None self.book_data = None self.embeddings = None def load_model(self, model_dir='model'): """Загрузка модели из сохраненных файлов""" try: # Загружаем данные книг self.book_data = pd.read_csv(f"{model_dir}/book_data.csv") # Загружаем эмбеддинги self.embeddings = np.load(f"{model_dir}/embeddings.npy") # Загружаем индекс FAISS self.faiss_index = faiss.read_index(f"{model_dir}/faiss_index.bin") return True except Exception as e: st.error(f"Ошибка при загрузке модели: {e}") return False def search(self, query, k=5): """Поиск книг по пользовательскому запросу""" if self.embedder is None or self.faiss_index is None: st.error("Поисковая система не инициализирована") return [] # Предобработка запроса processed_query = preprocess_text(query) # Получение эмбеддинга запроса query_embedding = self.embedder.get_embedding(processed_query) query_embedding = query_embedding.reshape(1, -1) # Нормализуем вектор запроса faiss.normalize_L2(query_embedding) # Поиск ближайших соседей scores, indices = self.faiss_index.search(query_embedding, k) # Формирование результатов results = [] for i, (score, idx) in enumerate(zip(scores[0], indices[0])): if idx < len(self.book_data): book = self.book_data.iloc[idx] results.append({ 'rank': i + 1, 'score': float(score), 'title': book.get('title', 'Нет названия'), 'author': book.get('author', 'Нет автора'), 'annotation': book.get('annotation', 'Нет аннотации'), 'page_url': book.get('page_url', '#'), 'book_image_url': book.get('book_image_url', book.get('image_url', '')) }) return results # Инициализация поисковой системы @st.cache_resource def initialize_search_engine(): # Инициализация модели RuBERT embedder = RuBERTEmbedder() # Создание поисковой системы search_engine = BookSearchEngine(embedder) # Загрузка подготовленной модели if search_engine.load_model(): st.success(f"Поисковая система загружена. Всего книг: {len(search_engine.book_data)}") else: st.error("Не удалось загрузить модель. Пожалуйста, убедитесь, что директория 'model' содержит необходимые файлы.") st.info("Перед запуском приложения нужно выполнить предварительную обработку данных с помощью скрипта preprocess.py") return search_engine # Основной код приложения def main(): st.set_page_config( page_title="Умный поиск книг", page_icon="📚", layout="wide" ) st.title("📚 Умный поиск книг") st.subheader("Найдите книги, соответствующие вашему запросу") # Инициализация поисковой системы search_engine = initialize_search_engine() # Основной интерфейс поиска st.write("### Введите описание книги, которую вы ищете") col1, col2 = st.columns([3, 1]) with col1: query = st.text_area("Описание книги:", height=150) with col2: num_results = st.slider("Количество результатов:", min_value=1, max_value=20, value=5) search_button = st.button("🔍 Искать", type="primary") # Если нажата кнопка поиска if search_button: if query: with st.spinner("Ищем подходящие книги..."): results = search_engine.search(query, k=num_results) if results: st.write(f"### Найдено {len(results)} подходящих книг:") for i, result in enumerate(results): col_image, col_content, col_score = st.columns([1, 2, 1]) with col_image: if 'book_image_url' in result and result['book_image_url']: try: st.image(result['book_image_url'], width=150) except Exception: st.write("Изображение недоступно") with col_content: if 'page_url' in result and result['page_url']: st.markdown(f"#### [{i+1}. {result['title']}]({result['page_url']})") else: st.markdown(f"#### {i+1}. {result['title']}") st.write(f"**Автор:** {result['author']}") with st.expander("Показать аннотацию"): st.write(result['annotation']) with col_score: st.metric( "Релевантность", f"{result['score']:.2f}", delta=None ) st.divider() else: st.info("К сожалению, подходящих книг не найдено.") else: st.warning("Пожалуйста, введите описание книги для поиска.") st.markdown("---") st.markdown("### О проекте") st.write(""" Этот сервис позволяет искать книги по их описанию с использованием семантической близости. Система анализирует смысл вашего запроса и находит книги с наиболее подходящими аннотациями. **Технологии:** - RuBERT для создания векторных представлений текста - FAISS для быстрого поиска ближайших соседей - Streamlit для веб-интерфейса """) if __name__ == "__main__": main()