import streamlit as st from datasets import load_dataset import numpy as np from sentence_transformers import SentenceTransformer import time from datetime import datetime import json import os import pandas as pd # =================================================================== # 1. НАСТРОЙКИ И ЗАГРУЗКА ДАННЫХ # =================================================================== MODEL_NAME = "BAAI/bge-large-en-v1.5" DATASET_NAME = "wikimedia/wikipedia" LANGUAGE = "20231101.en" ARTICLE_LIMIT = 1000 ADMIN_USER = "admin" ADMIN_PASS = "hfpassword21" LOG_FILE = "query_logs.json" # Настройка страницы st.set_page_config( page_title="Wikipedia Assistant", page_icon="🌖", layout="wide", initial_sidebar_state="expanded" ) # Загружаем логи if os.path.exists(LOG_FILE): with open(LOG_FILE, "r") as f: query_logs = json.load(f) else: query_logs = [] # =================================================================== # 2. ЗАГРУЗКА ДАННЫХ (с кэшированием) # =================================================================== @st.cache_resource def load_wikipedia(): """Загружает датасет Википедии""" with st.spinner("📚 Загружаю Википедию..."): dataset = load_dataset(DATASET_NAME, LANGUAGE, split="train", streaming=True) articles = [] for i, row in enumerate(dataset): if i >= ARTICLE_LIMIT: break articles.append({ "id": row.get("id", i), "title": row.get("title", "Без названия"), "text": row.get("text", "")[:3000], "url": row.get("url", "") }) return articles @st.cache_resource def load_embedder(): """Загружает модель для эмбеддингов""" with st.spinner("🧠 Загружаю модель..."): return SentenceTransformer(MODEL_NAME) @st.cache_resource def create_embeddings(articles, embedder): """Создаёт эмбеддинги статей""" with st.spinner("🔢 Создаю эмбеддинги статей..."): texts = [f"{a['title']}\n\n{a['text']}" for a in articles] return embedder.encode(texts, normalize_embeddings=True) # Загружаем всё articles = load_wikipedia() embedder = load_embedder() embeddings = create_embeddings(articles, embedder) # =================================================================== # 3. ФУНКЦИИ ПОИСКА И ЛОГИРОВАНИЯ # =================================================================== def search_wikipedia(query): """Ищет ответ на вопрос в Википедии""" if not query: return None start_time = time.time() # Поиск query_vector = embedder.encode([query], normalize_embeddings=True)[0] scores = embeddings @ query_vector top_indices = np.argsort(-scores)[:3] results = [] for idx in top_indices: score = float(scores[int(idx)]) if score > 0.3: article = articles[int(idx)] results.append({ "title": article['title'], "score": score, "text": article['text'][:1000], "url": article['url'] }) # Логируем запрос log_entry = { "timestamp": datetime.now().isoformat(), "query": query, "results_count": len(results), "response_time": round(time.time() - start_time, 2) } query_logs.append(log_entry) # Сохраняем логи with open(LOG_FILE, "w") as f: json.dump(query_logs[-100:], f) return results def get_admin_stats(): """Собирает статистику для админ-панели""" total_queries = len(query_logs) if total_queries > 0: avg_time = sum(q["response_time"] for q in query_logs) / total_queries popular_queries = sorted(query_logs, key=lambda x: x["results_count"], reverse=True)[:5] else: avg_time = 0 popular_queries = [] return { "total_queries": total_queries, "avg_time": avg_time, "popular_queries": popular_queries, "articles_count": len(articles), "model_name": MODEL_NAME } def clear_logs(): """Очищает логи""" global query_logs query_logs = [] with open(LOG_FILE, "w") as f: json.dump(query_logs, f) return True # =================================================================== # 4. ИНТЕРФЕЙС # =================================================================== # --- БОКОВАЯ ПАНЕЛЬ (АДМИНКА) --- with st.sidebar: st.image("https://cdn-icons-png.flaticon.com/512/4248/4248455.png", width=80) st.title("👑 Админ-панель") # Вход if "logged_in" not in st.session_state: st.session_state.logged_in = False if not st.session_state.logged_in: with st.form("login_form"): username = st.text_input("👤 Логин", placeholder="admin") password = st.text_input("🔑 Пароль", type="password", placeholder="hfpassword21") submitted = st.form_submit_button("🔑 Войти") if submitted: if username == ADMIN_USER and password == ADMIN_PASS: st.session_state.logged_in = True st.success("✅ Доступ разрешён!") st.rerun() else: st.error("❌ Неверный логин или пароль") else: st.success("✅ Вы вошли как администратор") # Кнопка выхода if st.button("🚪 Выйти"): st.session_state.logged_in = False st.rerun() st.divider() # Статистика st.subheader("📊 Статистика") stats = get_admin_stats() col1, col2 = st.columns(2) with col1: st.metric("Всего запросов", stats["total_queries"]) st.metric("Загружено статей", stats["articles_count"]) with col2: st.metric("Ср. время ответа", f"{stats['avg_time']:.2f}с") st.metric("Модель", stats["model_name"].split("/")[-1]) # Популярные запросы if stats["popular_queries"]: st.subheader("🔥 Топ-5 запросов") for i, q in enumerate(stats["popular_queries"], 1): st.write(f"{i}. **{q['query']}** (найдено: {q['results_count']})") st.divider() # Логи st.subheader("📋 Последние запросы") if query_logs: # Показываем последние 10 в виде таблицы df = pd.DataFrame(query_logs[-10:]) df["timestamp"] = pd.to_datetime(df["timestamp"]).dt.strftime("%H:%M:%S") st.dataframe( df[["timestamp", "query", "results_count", "response_time"]], column_config={ "timestamp": "Время", "query": "Запрос", "results_count": "Результатов", "response_time": "Время (с)" }, use_container_width=True, hide_index=True ) if st.button("🗑️ Очистить логи", type="secondary"): clear_logs() st.success("Логи очищены!") st.rerun() else: st.info("📭 Логов пока нет") # --- ОСНОВНАЯ ЧАСТЬ --- st.title("🌖 Wikipedia Assistant") st.markdown("Задай вопрос — я найду ответ в Википедии!") # Информация о загрузке st.success(f"✅ Загружено {len(articles)} статей") # Поиск query = st.text_input( "🔍 Что хочешь узнать?", placeholder="Например: Как измеряют расстояние до галактик?", key="query_input" ) col1, col2 = st.columns([1, 5]) with col1: search_clicked = st.button("🔎 Найти", type="primary", use_container_width=True) # Выполняем поиск if query and (search_clicked or query != st.session_state.get("last_query", "")): st.session_state.last_query = query with st.spinner("🔎 Ищу ответ..."): results = search_wikipedia(query) if results: for i, result in enumerate(results, 1): with st.expander(f"#{i} {result['title']} (сходство: {result['score']:.2f})", expanded=i==1): st.write(result['text'] + "...") if result['url']: st.link_button("🔗 Читать на Википедии", result['url']) else: st.warning("😕 Не нашёл подходящей статьи. Попробуй уточнить вопрос.") elif not query: st.info("💡 Напиши вопрос, и я найду ответ в Википедии") # =================================================================== # 5. ДОПОЛНИТЕЛЬНЫЕ ВОЗМОЖНОСТИ # =================================================================== # Подвал st.divider() st.caption("🌖 Wikipedia Assistant | Использует BAAI/bge-large-en-v1.5") # Кнопка внизу для переключения темы if st.button("🎨 Сменить тему"): st.session_state.theme = "dark" if st.session_state.get("theme") != "dark" else "light" st.rerun()