|
|
|
|
|
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'))
|
|
|
|
|
|
|
|
|
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()
|
|
|
|
|
|
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")
|
|
|
|
|
|
|
|
|
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():
|
|
|
|
|
|
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() |
