Update app.py

app.py

@@ -1,12 +1,10 @@
 #!/usr/bin/env python3
 """
-Демо RAG система для HuggingFace Spaces
-Использует предварительно обработанные чанки отчета Сбера
-Оптимизирована для быстрого запуска без тяжелых зависимостей
+Финальная векторная RAG система для HuggingFace Spaces
+Адаптированная версия с поддержкой векторного поиска и резервным режимом
 """
 
 import os
-import sys
 import json
 import pickle
 import tempfile
@@ -15,157 +13,276 @@ from typing import Optional, Dict, Any, List, Tuple
 import traceback
 import re
 
-import gradio as gr
-import numpy as np
+try:
+    import numpy as np
+    import faiss
+    HAS_FAISS = True
+except ImportError:
+    HAS_FAISS = False
+    print("⚠️ FAISS не установлен, будет использован поиск по ключевым словам")
+
+try:
+    import gradio as gr
+    HAS_GRADIO = True
+except ImportError:
+    HAS_GRADIO = False
+    print("⚠️ Gradio не установлен")
 
-# OpenAI для генерации ответов
 from openai import OpenAI
 
-class LightweightRAGSystem:
-    """Легковесная RAG система с предзагруженными чанками"""
+class VectorRAGSystem:
+    """RAG система с векторным поиском и резервным режимом"""
     
     def __init__(self):
         self.chunks = []
         self.word_index = {}
+        self.faiss_index = None
         self.metadata = {}
         self.client = None
         self.is_initialized = False
         
-        # Конфигурация
+        # Модели и параметры
+        self.embedding_model = "text-embedding-3-large"
+        self.embedding_dim = 3072
         self.generation_model = "gpt-4o"
         self.reranking_model = "gpt-4o-mini"
+        
+        # Параметры поиска
         self.max_chunks_for_rerank = 15
         self.final_chunks_count = 5
+        self.vector_search_k = 20
+        
+        # Режим работы
+        self.vector_mode = HAS_FAISS
     
-    def load_preprocessed_data(self) -> bool:
-        """Загрузка предварительно обработанных данных"""
+    def initialize_with_api_key(self, api_key: str) -> Tuple[str, str]:
+        """Инициализация системы с API ключом"""
         try:
-            print("🔄 Загрузка предварительно обработанных данных...")
+            if not api_key.strip():
+                return "❌ Введите OpenAI API ключ", ""
             
-            # Загружаем улучшенный индекс с таблицами
-            index_file = "enhanced_sber_index.pkl"
-            if not os.path.exists(index_file):
-                print(f"❌ Файл индекса не найден: {index_file}")
+            # Инициализация OpenAI клиента
+            self.client = OpenAI(api_key=api_key.strip())
+            
+            # Загрузка данных
+            if not self.load_data():
+                return "❌ Ошибка загрузки данных", ""
+            
+            self.is_initialized = True
+            stats = self._generate_stats()
+            
+            return "✅ Векторная RAG система инициализирована", stats
+            
+        except Exception as e:
+            return f"❌ Ошибка инициализации: {str(e)}", ""
+    
+    def load_data(self) -> bool:
+        """Загрузка данных (векторных или обычных)"""
+        try:
+            # Сначала пробуем загрузить векторные данные
+            if self.vector_mode and self.load_vector_data():
+                return True
+            
+            # Если не удалось, загружаем обычные данные
+            return self.load_fallback_data()
+            
+        except Exception as e:
+            print(f"❌ Ошибка загрузки данных: {e}")
+            return False
+    
+    def load_vector_data(self) -> bool:
+        """Загрузка векторных данных"""
+        try:
+            print("🔄 Попытка загрузки векторных данных...")
+            
+            # Файлы векторных данных
+            chunks_file = "vector_enhanced_sber_chunks.pkl"
+            metadata_file = "vector_enhanced_sber_metadata.json"
+            faiss_file = "vector_enhanced_sber_faiss.index"
+            
+            if not all(os.path.exists(f) for f in [chunks_file, metadata_file, faiss_file]):
+                print("📁 Файлы векторных данных не найдены")
                 return False
             
-            with open(index_file, 'rb') as f:
-                index_data = pickle.load(f)
+            # Загружаем чанки
+            with open(chunks_file, 'rb') as f:
+                chunks_data = pickle.load(f)
+            
+            self.chunks = []
+            for chunk_data in chunks_data:
+                self.chunks.append({
+                    "text": chunk_data["text"],
+                    "page": chunk_data["page"],
+                    "chunk_index": chunk_data["chunk_index"],
+                    "embedding": np.array(chunk_data["embedding"]) if chunk_data.get("embedding") else None,
+                    "metadata": chunk_data.get("metadata", {}),
+                    "full_page_text": chunk_data.get("full_page_text", chunk_data["text"])
+                })
             
-            self.chunks = index_data["chunks"]
-            self.word_index = index_data["word_index"]
-            self.metadata = index_data["metadata"]
+            # Загружаем метаданные
+            with open(metadata_file, 'r', encoding='utf-8') as f:
+                self.metadata = json.load(f)
             
-            print(f"✅ Загружено {len(self.chunks)} чанков")
-            print(f"✅ Создан словарный индекс из {len(self.word_index)} слов")
+            # Загружаем FAISS индекс
+            if HAS_FAISS:
+                self.faiss_index = faiss.read_index(faiss_file)
             
+            print(f"✅ Загружены векторные данные: {len(self.chunks)} чанков")
             return True
             
         except Exception as e:
-            print(f"❌ Ошибка загрузки данных: {e}")
-            traceback.print_exc()
+            print(f"❌ Ошибка загрузки векторных данных: {e}")
             return False
     
-    def initialize_with_api_key(self, api_key: str) -> Tuple[str, str]:
-        """Инициализация системы с API ключом"""
+    def load_fallback_data(self) -> bool:
+        """Загрузка обычных данных"""
         try:
-            if not api_key.strip():
-                return "❌ Введите OpenAI API ключ", ""
+            print("🔄 Загрузка резервных данных...")
             
-            # Инициализация OpenAI клиента
-            self.client = OpenAI(api_key=api_key.strip())
+            index_file = "enhanced_sber_index.pkl"
+            if not os.path.exists(index_file):
+                print(f"❌ Файл резервных данных не найден: {index_file}")
+                return False
             
-            # Загрузка данных
-            if not self.load_preprocessed_data():
-                return "❌ Ошибка загрузки данных", ""
+            with open(index_file, 'rb') as f:
+                index_data = pickle.load(f)
             
-            self.is_initialized = True
+            # Конвертируем в формат чанков
+            self.chunks = []
+            chunk_texts = index_data.get("chunks", [])
+            
+            for i, chunk_text in enumerate(chunk_texts):
+                chunk = {
+                    "text": chunk_text,
+                    "page": index_data.get("metadata", {}).get("chunk_pages", {}).get(str(i), 1),
+                    "chunk_index": i,
+                    "embedding": None,
+                    "metadata": {},
+                    "full_page_text": chunk_text
+                }
+                self.chunks.append(chunk)
             
-            # Генерация статистики
-            stats = self._generate_stats()
+            # Создаем словарный индекс для поиска
+            self.word_index = index_data.get("word_index", {})
+            self.metadata = index_data.get("metadata", {})
             
-            return "✅ Система инициализирована успешно", stats
+            self.vector_mode = False  # Отключаем векторный режим
+            
+            print(f"✅ Загружены резервные данные: {len(self.chunks)} чанков")
+            return True
             
         except Exception as e:
-            return f"❌ Ошибка инициализации: {str(e)}", ""
+            print(f"❌ Ошибка загрузки резервных данных: {e}")
+            return False
     
     def _generate_stats(self) -> str:
         """Генерация статистики системы"""
-        total_chunks = self.metadata.get("total_chunks", 0)
-        avg_length = self.metadata.get("avg_chunk_length", 0)
-        avg_tokens = self.metadata.get("avg_token_count", 0)
-        pages = self.metadata.get("pages_processed", 0)
-        
-        # Добавим информацию о таблицах
-        text_chunks = self.metadata.get("text_chunks", 0)
-        table_chunks = self.metadata.get("table_chunks", 0)
-        table_pages = self.metadata.get("table_pages", 0)
+        total_chunks = len(self.chunks)
+        mode = "Векторный поиск" if self.vector_mode and self.faiss_index else "Поиск по ключевым словам"
         
-        stats = f"""✅ **Улучшенная система готова к работе!**
+        stats = f"""✅ **RAG система готова!**
 
 📊 **Статистика:**
 - 📦 Загружено чанков: {total_chunks}
-- 📝 Текстовых чанков: {text_chunks}
-- 📋 Табличных чанков: {table_chunks}
-- 📏 Средняя длина чанка: {avg_length:.0f} символов
-- 🔢 Средний размер: {avg_tokens:.0f} токенов
-- 📖 Страниц отчета: {pages}
-- 📊 Страниц с таблицами: {table_pages}
+- 🔍 Режим поиска: {mode}
+- 🧠 Модель генерации: {self.generation_model}
+- 🎯 Реранкинг: {self.reranking_model}
 
 🔍 **Возможности:**
-- 🔎 Быстрый поиск по ключевым словам
-- 📋 Извлечение структурированных таблиц
-- 🧠 LLM реранкинг результатов (GPT-4o-mini)
-- 📝 Интеллектуальная генерация ответов (GPT-4o)
+- 🔎 Семантический/ключевой поиск
+- 📄 Контекстное обогащение
+- 🧠 LLM реранкинг результатов
+- 📝 Интеллектуальная генерация ответов
 - 📊 Анализ годового отчета ПАО Сбербанк 2023
 
-🚀 **Готова отвечать на вопросы с поддержкой таблиц!**"""
+🚀 **Готова к работе!**"""
         
         return stats
     
-    def search_by_keywords(self, query: str, max_results: int = 30) -> List[Dict]:
-        """Поиск по ключевым словам"""
-        if not query.strip():
-            return []
+    def search(self, query: str, k: int = 20) -> List[Tuple[Dict, float]]:
+        """Основной метод поиска"""
+        if self.vector_mode and self.faiss_index and self.client:
+            return self.vector_search(query, k)
+        else:
+            return self.keyword_search(query, k)
+    
+    def vector_search(self, query: str, k: int = 20) -> List[Tuple[Dict, float]]:
+        """Векторный поиск по запросу"""
+        if not self.faiss_index or not self.client:
+            return self.keyword_search(query, k)
         
-        # Извлекаем ключевые слова из запроса
+        try:
+            # Создаем эмбеддинг для запроса
+            response = self.client.embeddings.create(
+                model=self.embedding_model,
+                input=[query]
+            )
+            
+            query_embedding = np.array(response.data[0].embedding, dtype=np.float32)
+            query_embedding = query_embedding.reshape(1, -1)
+            
+            # Нормализуем для Inner Product
+            faiss.normalize_L2(query_embedding)
+            
+            # Поиск в FAISS индексе
+            scores, indices = self.faiss_index.search(query_embedding, k)
+            
+            # Формируем результаты
+            results = []
+            for score, idx in zip(scores[0], indices[0]):
+                if 0 <= idx < len(self.chunks):
+                    chunk = self.chunks[idx]
+                    results.append((chunk, float(score)))
+            
+            return results
+            
+        except Exception as e:
+            print(f"❌ Ошибка векторного поиска: {e}")
+            return self.keyword_search(query, k)
+    
+    def keyword_search(self, query: str, k: int = 20) -> List[Tuple[Dict, float]]:
+        """Поиск по ключевым словам"""
         query_words = set(re.findall(r'\b\w+\b', query.lower()))
         
-        # Находим чанки, содержащие эти слова
-        chunk_scores = {}
-        
-        for word in query_words:
-            if word in self.word_index:
-                for chunk_idx in self.word_index[word]:
-                    if chunk_idx not in chunk_scores:
-                        chunk_scores[chunk_idx] = 0
-                    chunk_scores[chunk_idx] += 1
+        if self.word_index:
+            # Используем готовый индекс
+            chunk_scores = {}
+            for word in query_words:
+                if word in self.word_index:
+                    for chunk_idx in self.word_index[word]:
+                        if chunk_idx not in chunk_scores:
+                            chunk_scores[chunk_idx] = 0
+                        chunk_scores[chunk_idx] += 1
+        else:
+            # Создаем индекс на лету
+            chunk_scores = {}
+            for i, chunk in enumerate(self.chunks):
+                text_words = set(re.findall(r'\b\w+\b', chunk["text"].lower()))
+                score = len(query_words.intersection(text_words))
+                if score > 0:
+                    chunk_scores[i] = score
         
-        # Сортируем по количеству совпадений
+        # Сортируем по скору
         sorted_chunks = sorted(chunk_scores.items(), key=lambda x: x[1], reverse=True)
         
-        # Возвращаем результаты
         results = []
-        for chunk_idx, score in sorted_chunks[:max_results]:
+        for chunk_idx, score in sorted_chunks[:k]:
             if chunk_idx < len(self.chunks):
-                chunk = self.chunks[chunk_idx].copy()
-                chunk["keyword_score"] = score
-                chunk["similarity"] = score / len(query_words)  # Нормализованный score
-                results.append(chunk)
+                chunk = self.chunks[chunk_idx]
+                results.append((chunk, float(score)))
         
         return results
     
-    def rerank_with_llm(self, query: str, chunks: List[Dict]) -> List[Dict]:
+    def rerank_with_llm(self, query: str, chunks: List[Tuple[Dict, float]]) -> List[Tuple[Dict, float]]:
         """LLM реранкинг результатов"""
         if not chunks or not self.client:
             return chunks
         
         try:
-            # Ограничиваем количество чанков для реранкинга
             chunks_to_rerank = chunks[:self.max_chunks_for_rerank]
             
-            # Подготавливаем документы для реранкинга
             docs_text = ""
-            for i, chunk in enumerate(chunks_to_rerank):
+            for i, (chunk, _) in enumerate(chunks_to_rerank):
                 preview = chunk['text'][:300] + "..." if len(chunk['text']) > 300 else chunk['text']
                 docs_text += f"\nДокумент {i+1} (стр. {chunk['page']}):\n{preview}\n"
             
@@ -190,51 +307,35 @@ class LightweightRAGSystem:
                 temperature=0
             )
             
-            # Парсим оценки
             scores_text = response.choices[0].message.content.strip()
-            scores = []
-            
             numbers = re.findall(r'\d+\.?\d*', scores_text)
-            for num in numbers:
-                score = float(num)
-                score = max(0, min(10, score))  # Ограничиваем 0-10
-                scores.append(score)
+            scores = [max(0, min(10, float(num))) for num in numbers]
             
-            # Применяем оценки
             reranked = []
-            for i, chunk in enumerate(chunks):
-                chunk_copy = chunk.copy()
-                if i < len(scores):
-                    chunk_copy["rerank_score"] = scores[i]
-                else:
-                    chunk_copy["rerank_score"] = 0
-                reranked.append(chunk_copy)
-            
-            # Сортируем по реранк скору
-            reranked.sort(key=lambda x: x["rerank_score"], reverse=True)
+            for i, (chunk, original_score) in enumerate(chunks):
+                rerank_score = scores[i] if i < len(scores) else 0
+                reranked.append((chunk, rerank_score))
             
+            reranked.sort(key=lambda x: x[1], reverse=True)
             return reranked
             
         except Exception as e:
             print(f"❌ Ошибка реранкинга: {e}")
             return chunks
     
-    def generate_answer(self, query: str, context_chunks: List[Dict]) -> str:
+    def generate_answer(self, query: str, context_chunks: List[Tuple[Dict, float]]) -> str:
         """Генерация ответа на основе контекста"""
         if not self.client:
             return "❌ OpenAI API не настроен"
         
         try:
-            # Подготавливаем контекст
             context_parts = []
-            for i, chunk in enumerate(context_chunks[:self.final_chunks_count]):
-                # Очищаем текст от проблемных символов
-                clean_text = chunk['text'].encode('utf-8', errors='ignore').decode('utf-8')
-                context_parts.append(f"Фрагмент {i+1} (страница {chunk['page']}):\n{clean_text}")
+            for i, (chunk, score) in enumerate(context_chunks[:self.final_chunks_count]):
+                text = chunk.get('full_page_text', chunk['text'])
+                clean_text = text.encode('utf-8', errors='ignore').decode('utf-8')
+                context_parts.append(f"Фрагмент {i+1} (страница {chunk['page']}, релевантность: {score:.2f}):\n{clean_text}")
             
             context = "\n\n".join(context_parts)
-            
-            # Очищаем запрос
             clean_query = query.encode('utf-8', errors='ignore').decode('utf-8')
             
             prompt = f"""Ты - эк��перт по анализу финансовых отчетов. Ответь на вопрос пользователя на основе предоставленного контекста из годового отчета ПАО Сбербанк 2023.
@@ -261,11 +362,7 @@ class LightweightRAGSystem:
                 temperature=0.1
             )
             
-            answer = response.choices[0].message.content
-            if answer:
-                return answer.strip()
-            else:
-                return "Получен пустой ответ от модели"
+            return response.choices[0].message.content.strip()
             
         except Exception as e:
             return f"❌ Ошибка генерации ответа: {str(e)}"
@@ -287,39 +384,37 @@ class LightweightRAGSystem:
             }
         
         try:
-            # Шаг 1: Поиск по ключевым словам
-            initial_results = self.search_by_keywords(query, max_results=30)
+            # Поиск
+            search_results = self.search(query, k=self.vector_search_k)
             
-            if not initial_results:
+            if not search_results:
                 return {
                     "answer": "К сожалению, не удалось найти релевантную информацию по вашему вопросу.",
                     "sources": [],
-                    "debug_info": {"step": "keyword_search", "results_count": 0}
+                    "debug_info": {"step": "search", "results_count": 0}
                 }
             
-            # Шаг 2: LLM реранкинг
-            reranked_results = self.rerank_with_llm(query, initial_results)
+            # Реранкинг
+            reranked_results = self.rerank_with_llm(query, search_results)
             
-            # Шаг 3: Генерация ответа
-            top_chunks = reranked_results[:self.final_chunks_count]
-            answer = self.generate_answer(query, top_chunks)
+            # Генерация ответа
+            answer = self.generate_answer(query, reranked_results)
             
             # Подготовка источников
             sources = []
-            for chunk in top_chunks:
+            for chunk, score in reranked_results[:self.final_chunks_count]:
                 sources.append({
                     "page": chunk["page"],
-                    "keyword_score": chunk.get("keyword_score", 0),
-                    "rerank_score": chunk.get("rerank_score", 0),
+                    "search_score": search_results[0][1] if search_results else 0,
+                    "rerank_score": score,
                     "preview": chunk["text"][:200] + "..." if len(chunk["text"]) > 200 else chunk["text"]
                 })
             
             debug_info = {
-                "initial_results": len(initial_results),
+                "search_results": len(search_results),
                 "reranked_results": len(reranked_results),
-                "final_chunks": len(top_chunks),
-                "avg_keyword_score": np.mean([s["keyword_score"] for s in sources]) if sources else 0,
-                "avg_rerank_score": np.mean([s["rerank_score"] for s in sources]) if sources else 0
+                "final_chunks": len(sources),
+                "search_method": "vector" if self.vector_mode else "keyword"
             }
             
             return {
@@ -338,7 +433,7 @@ class LightweightRAGSystem:
             }
 
 # Глобальная переменная системы
-rag_system = LightweightRAGSystem()
+rag_system = VectorRAGSystem()
 
 def initialize_system(api_key: str) -> Tuple[str, str]:
     """Инициализация системы"""
@@ -356,7 +451,7 @@ def ask_question(question: str) -> Tuple[str, str]:
         sources_info = "\n📚 **Источники:**\n"
         for i, source in enumerate(result["sources"], 1):
             sources_info += f"\n**{i}.** Страница {source['page']} "
-            sources_info += f"(ключевые слова: {source['keyword_score']}, "
+            sources_info += f"(поиск: {source['search_score']:.3f}, "
             sources_info += f"релевантность: {source['rerank_score']:.1f}/10)\n"
             sources_info += f"*Превью:* {source['preview']}\n"
     
@@ -364,19 +459,22 @@ def ask_question(question: str) -> Tuple[str, str]:
     if result.get("debug_info"):
         debug = result["debug_info"]
         sources_info += f"\n🔍 **Статистика поиска:**\n"
-        sources_info += f"- Найдено по ключевым словам: {debug.get('initial_results', 0)}\n"
+        sources_info += f"- Метод поиска: {debug.get('search_method', 'unknown')}\n"
+        sources_info += f"- Найдено результатов: {debug.get('search_results', 0)}\n"
         sources_info += f"- После реранкинга: {debug.get('reranked_results', 0)}\n"
         sources_info += f"- Использовано в отв��те: {debug.get('final_chunks', 0)}\n"
-        if debug.get('avg_rerank_score'):
-            sources_info += f"- Средняя релевантность: {debug.get('avg_rerank_score', 0):.1f}/10\n"
     
     return answer, sources_info
 
 def create_demo_interface():
-    """Создание демо интерфейса для HF"""
+    """Создание демо интерфейса"""
+    
+    if not HAS_GRADIO:
+        print("❌ Gradio не установлен. Установите: pip install gradio")
+        return None
     
     with gr.Blocks(
-        title="RAG Demo - Сбер 2023",
+        title="Vector RAG Demo - Сбер 2023",
         theme=gr.themes.Soft(),
         css="""
         .main-header { text-align: center; margin-bottom: 2rem; }
@@ -386,9 +484,9 @@ def create_demo_interface():
         
         gr.Markdown("""
         <div class="main-header">
-        <h1>🏆 Enhanced RAG Demo: Анализ отчета Сбера 2023</h1>
-        <p>Улучшенная система поиска с поддержкой таблиц</p>
-        <p><strong>84 извлеченные таблицы • 2009 чанков • pdfplumber обработка</strong></p>
+        <h1>🚀 Advanced RAG Demo: Анализ отчета Сбера 2023</h1>
+        <p>Умная система с векторным поиском и адаптивным режимом</p>
+        <p><strong>OpenAI embeddings • FAISS IndexFlatIP • LLM reranking • Fallback mode</strong></p>
         </div>
         """)
         
@@ -398,9 +496,9 @@ def create_demo_interface():
                 
                 api_key_input = gr.Textbox(
                     label="OpenAI API Key",
-                    placeholder="sk-...",
+                    placeholder="sk-proj-...",
                     type="password",
-                    info="Введите ваш OpenAI API ключ для работы системы"
+                    info="Введите ваш OpenAI API ключ"
                 )
                 
                 init_btn = gr.Button("🚀 Инициализировать", variant="primary")
@@ -423,7 +521,7 @@ def create_demo_interface():
                 lines=2,
                 scale=4
             )
-            ask_btn = gr.Button("📝 Спросить", variant="primary", scale=1)
+            ask_btn = gr.Button("🔍 Поиск", variant="primary", scale=1)
         
         with gr.Row():
             with gr.Column(scale=2):
@@ -447,6 +545,7 @@ def create_demo_interface():
         - Расскажите о кредитном портфеле Сбербанка
         - Какие технологические инициативы развивает Сбер?
         - Каковы показатели рентабельности банка?
+        - Какие ESG инициативы реализует Сбер?
         """)
         
         # Event handlers
@@ -470,11 +569,6 @@ def create_demo_interface():
     
     return demo
 
-if __name__ == "__main__":
-    demo = create_demo_interface()
-    demo.launch(
-        share=False,
-        server_name="0.0.0.0",
-        server_port=7860,
-        show_error=True
-    )
+# Запуск для Hugging Face Spaces
+demo = create_demo_interface()
+demo.launch()