app.py

# 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()