fruitpicker01 commited on
Commit
d9d149a
·
verified ·
1 Parent(s): e04b8e2

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +52 -88
app.py CHANGED
@@ -40,6 +40,8 @@ class VectorRAGSystem:
40
  self.metadata = {}
41
  self.client = None
42
  self.is_initialized = False
 
 
43
 
44
  # Модели и параметры
45
  self.embedding_model = "text-embedding-3-large"
@@ -77,14 +79,14 @@ class VectorRAGSystem:
77
  return f"❌ Ошибка инициализации: {str(e)}", ""
78
 
79
  def load_data(self) -> bool:
80
- """Загрузка данных (векторных или обычных)"""
81
  try:
82
- # Сначала пробуем загрузить векторные данные
83
  if self.vector_mode and self.load_vector_data():
84
  return True
85
 
86
- # Если не удалось, загружаем обычные данные
87
- return self.load_fallback_data()
88
 
89
  except Exception as e:
90
  print(f"❌ Ошибка загрузки данных: {e}")
@@ -103,28 +105,38 @@ class VectorRAGSystem:
103
  print("📁 Файлы векторных данных не найдены")
104
  return False
105
 
106
- # Загружаем метаданные с чанками
107
  with open(metadata_file, 'r', encoding='utf-8') as f:
108
- full_data = json.load(f)
109
 
110
- # Извлекаем чанки
111
- chunks_data = full_data.get("chunks", [])
112
  self.chunks = []
113
- for chunk_data in chunks_data:
114
  self.chunks.append({
115
- "text": chunk_data["text"],
116
- "page": chunk_data["page"],
117
- "chunk_index": chunk_data.get("chunk_index", len(self.chunks)),
118
- "metadata": chunk_data.get("metadata", {})
 
119
  })
120
 
121
- # Сохраняем остальные метаданные
122
- self.metadata = full_data
123
 
124
  # Загружаем FAISS индекс
125
  if HAS_FAISS:
126
  self.faiss_index = faiss.read_index(faiss_file)
127
 
 
 
 
 
 
 
 
 
 
 
128
  print(f"✅ Загружены векторные данные: {len(self.chunks)} чанков")
129
  return True
130
 
@@ -132,46 +144,25 @@ class VectorRAGSystem:
132
  print(f"❌ Ошибка загрузки векторных данных: {e}")
133
  return False
134
 
135
- def load_fallback_data(self) -> bool:
136
- """Загрузка обычных данных"""
 
 
 
137
  try:
138
- print("🔄 Загрузка резервных данных...")
139
-
140
- index_file = "enhanced_sber_index.pkl"
141
- if not os.path.exists(index_file):
142
- print(f"❌ Файл резервных данных не найден: {index_file}")
143
- return False
144
-
145
- with open(index_file, 'rb') as f:
146
- index_data = pickle.load(f)
147
-
148
- # Конвертируем в формат чанков
149
- self.chunks = []
150
- chunk_texts = index_data.get("chunks", [])
151
-
152
- for i, chunk_text in enumerate(chunk_texts):
153
- chunk = {
154
- "text": chunk_text,
155
- "page": index_data.get("metadata", {}).get("chunk_pages", {}).get(str(i), 1),
156
- "chunk_index": i,
157
- "embedding": None,
158
- "metadata": {},
159
- "full_page_text": chunk_text
160
- }
161
- self.chunks.append(chunk)
162
-
163
- # Создаем словарный индекс для поиска
164
- self.word_index = index_data.get("word_index", {})
165
- self.metadata = index_data.get("metadata", {})
166
 
167
- self.vector_mode = False # Отключаем векторный режим
 
168
 
169
- print(f"✅ Загружены резервные данные: {len(self.chunks)} чанков")
170
- return True
 
171
 
172
  except Exception as e:
173
- print(f"❌ Ошибка загрузки резервных данных: {e}")
174
- return False
175
 
176
  def _generate_stats(self) -> str:
177
  """Генерация статистики системы"""
@@ -202,12 +193,14 @@ class VectorRAGSystem:
202
  if self.vector_mode and self.faiss_index and self.client:
203
  return self.vector_search(query, k)
204
  else:
205
- return self.keyword_search(query, k)
 
206
 
207
  def vector_search(self, query: str, k: int = 20) -> List[Tuple[Dict, float]]:
208
  """Векторный поиск по запросу"""
209
  if not self.faiss_index or not self.client:
210
- return self.keyword_search(query, k)
 
211
 
212
  try:
213
  # Создаем эмбеддинг для запроса
@@ -225,51 +218,22 @@ class VectorRAGSystem:
225
  # Поиск в FAISS индексе
226
  scores, indices = self.faiss_index.search(query_embedding, k)
227
 
228
- # Формируем результаты
229
  results = []
230
  for score, idx in zip(scores[0], indices[0]):
231
  if 0 <= idx < len(self.chunks):
232
- chunk = self.chunks[idx]
 
 
 
233
  results.append((chunk, float(score)))
234
 
235
  return results
236
 
237
  except Exception as e:
238
  print(f"❌ Ошибка векторного поиска: {e}")
239
- return self.keyword_search(query, k)
240
-
241
- def keyword_search(self, query: str, k: int = 20) -> List[Tuple[Dict, float]]:
242
- """Поиск по ключевым словам"""
243
- query_words = set(re.findall(r'\b\w+\b', query.lower()))
244
-
245
- if self.word_index:
246
- # Используем готовый индекс
247
- chunk_scores = {}
248
- for word in query_words:
249
- if word in self.word_index:
250
- for chunk_idx in self.word_index[word]:
251
- if chunk_idx not in chunk_scores:
252
- chunk_scores[chunk_idx] = 0
253
- chunk_scores[chunk_idx] += 1
254
- else:
255
- # Создаем индекс на лету
256
- chunk_scores = {}
257
- for i, chunk in enumerate(self.chunks):
258
- text_words = set(re.findall(r'\b\w+\b', chunk["text"].lower()))
259
- score = len(query_words.intersection(text_words))
260
- if score > 0:
261
- chunk_scores[i] = score
262
-
263
- # Сортируем по скору
264
- sorted_chunks = sorted(chunk_scores.items(), key=lambda x: x[1], reverse=True)
265
-
266
- results = []
267
- for chunk_idx, score in sorted_chunks[:k]:
268
- if chunk_idx < len(self.chunks):
269
- chunk = self.chunks[chunk_idx]
270
- results.append((chunk, float(score)))
271
-
272
- return results
273
 
274
  def rerank_with_llm(self, query: str, chunks: List[Tuple[Dict, float]]) -> List[Tuple[Dict, float]]:
275
  """LLM реранкинг результатов"""
 
40
  self.metadata = {}
41
  self.client = None
42
  self.is_initialized = False
43
+ self.pdf_doc = None
44
+ self.page_texts = {} # Кеш текстов страниц
45
 
46
  # Модели и параметры
47
  self.embedding_model = "text-embedding-3-large"
 
79
  return f"❌ Ошибка инициализации: {str(e)}", ""
80
 
81
  def load_data(self) -> bool:
82
+ """Загрузка векторных данных"""
83
  try:
84
+ # Загружаем только векторные данные
85
  if self.vector_mode and self.load_vector_data():
86
  return True
87
 
88
+ print("❌ Векторные данные не найдены или не удалось загрузить")
89
+ return False
90
 
91
  except Exception as e:
92
  print(f"❌ Ошибка загрузки данных: {e}")
 
105
  print("📁 Файлы векторных данных не найдены")
106
  return False
107
 
108
+ # Загружаем метаданные (список с page и chunk_id)
109
  with open(metadata_file, 'r', encoding='utf-8') as f:
110
+ metadata_list = json.load(f)
111
 
112
+ # Создаем структуру чанков без текстов (получим из PDF по требованию)
 
113
  self.chunks = []
114
+ for i, item in enumerate(metadata_list):
115
  self.chunks.append({
116
+ "page": item["page"],
117
+ "chunk_id": item["chunk_id"],
118
+ "chunk_index": i,
119
+ "text": "", # Получим из PDF по требованию
120
+ "metadata": {}
121
  })
122
 
123
+ # Сохраняем метаданные
124
+ self.metadata = {"total_chunks": len(self.chunks)}
125
 
126
  # Загружаем FAISS индекс
127
  if HAS_FAISS:
128
  self.faiss_index = faiss.read_index(faiss_file)
129
 
130
+ # Загружаем PDF для parent-page enrichment
131
+ pdf_path = "data/Сбер 2023.pdf"
132
+ if os.path.exists(pdf_path):
133
+ import fitz # PyMuPDF
134
+ self.pdf_doc = fitz.open(pdf_path)
135
+ print(f"✅ PDF загружен: {self.pdf_doc.page_count} страниц")
136
+ else:
137
+ print("❌ PDF файл не найден для parent-page enrichment")
138
+ self.pdf_doc = None
139
+
140
  print(f"✅ Загружены векторные данные: {len(self.chunks)} чанков")
141
  return True
142
 
 
144
  print(f"❌ Ошибка загрузки векторных данных: {e}")
145
  return False
146
 
147
+ def get_page_text(self, page_num: int) -> str:
148
+ """Получение полного текста страницы с кешированием"""
149
+ if page_num in self.page_texts:
150
+ return self.page_texts[page_num]
151
+
152
  try:
153
+ if not self.pdf_doc or page_num < 1 or page_num > self.pdf_doc.page_count:
154
+ return ""
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
155
 
156
+ page = self.pdf_doc[page_num - 1] # PyMuPDF использует 0-based индексы
157
+ text = page.get_text()
158
 
159
+ # Кешируем текст
160
+ self.page_texts[page_num] = text
161
+ return text
162
 
163
  except Exception as e:
164
+ print(f"❌ Ошибка получения текста страницы {page_num}: {e}")
165
+ return ""
166
 
167
  def _generate_stats(self) -> str:
168
  """Генерация статистики системы"""
 
193
  if self.vector_mode and self.faiss_index and self.client:
194
  return self.vector_search(query, k)
195
  else:
196
+ print("⚠️ Векторный режим отключен")
197
+ return []
198
 
199
  def vector_search(self, query: str, k: int = 20) -> List[Tuple[Dict, float]]:
200
  """Векторный поиск по запросу"""
201
  if not self.faiss_index or not self.client:
202
+ print("⚠️ FAISS индекс или OpenAI клиент недоступны")
203
+ return []
204
 
205
  try:
206
  # Создаем эмбеддинг для запроса
 
218
  # Поиск в FAISS индексе
219
  scores, indices = self.faiss_index.search(query_embedding, k)
220
 
221
+ # Формируем результаты с parent-page enrichment
222
  results = []
223
  for score, idx in zip(scores[0], indices[0]):
224
  if 0 <= idx < len(self.chunks):
225
+ chunk = self.chunks[idx].copy()
226
+ # Получаем полный текст страницы для parent-page enrichment
227
+ page_text = self.get_page_text(chunk["page"])
228
+ chunk["text"] = page_text if page_text else chunk["text"]
229
  results.append((chunk, float(score)))
230
 
231
  return results
232
 
233
  except Exception as e:
234
  print(f"❌ Ошибка векторного поиска: {e}")
235
+ print("⚠️ Переход на поиск без векторов невозможен")
236
+ return []
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
237
 
238
  def rerank_with_llm(self, query: str, chunks: List[Tuple[Dict, float]]) -> List[Tuple[Dict, float]]:
239
  """LLM реранкинг результатов"""