fruitpicker01 commited on
Commit
8cc0322
·
verified ·
1 Parent(s): d54d97c

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +35 -47
app.py CHANGED
@@ -119,60 +119,53 @@ class VectorRAGSystem:
119
  return False
120
 
121
  def load_vector_data(self) -> bool:
122
- """Загрузка векторных данных"""
123
  try:
124
  print("🔄 Попытка загрузки векторных данных...")
125
 
126
- # Файлы векторных данных
127
- faiss_file = "chunks_flatip.faiss"
128
- metadata_file = "metadata.json"
129
-
130
  if not all(os.path.exists(f) for f in [faiss_file, metadata_file]):
131
- print("📁 Файлы векторных данных не найдены")
132
  return False
133
 
134
- # Загружаем метаданные
135
  with open(metadata_file, 'r', encoding='utf-8') as f:
136
  metadata_list = json.load(f)
 
137
 
138
- # Собираем self.chunks для всех типов элементов
139
  self.chunks = []
140
  for i, item in enumerate(metadata_list):
141
- # унифицируем идентификатор чанка
142
  chunk_id = item.get("chunk_id",
143
  item.get("table_id",
144
  item.get("img_id", None)))
145
-
146
  self.chunks.append({
147
  "page": item["page"],
148
  "chunk_id": chunk_id,
149
  "chunk_index": i,
150
- "text": "", # подгрузим из PDF при выдаче
151
- "metadata": {}
152
  })
153
 
154
- # Сохраняем общую статистику
155
- self.metadata = {"total_chunks": len(self.chunks)}
156
-
157
- # Загружаем FAISS-индекс
158
  if HAS_FAISS:
159
  self.faiss_index = faiss.read_index(faiss_file)
160
 
161
- # Загружаем PDF для parent-page enrichment
162
  pdf_path = "data/Сбер 2023.pdf"
163
  if os.path.exists(pdf_path):
164
- import fitz # PyMuPDF
165
  self.pdf_doc = fitz.open(pdf_path)
166
  print(f"✅ PDF загружен: {self.pdf_doc.page_count} страниц")
167
  else:
168
- print("❌ PDF файл не найден для parent-page enrichment")
169
  self.pdf_doc = None
170
 
171
- print(f"✅ Загружены векторные данные: {len(self.chunks)} чанков")
172
  return True
173
-
174
  except Exception as e:
175
- print(f"❌ Ошибка загрузки векторных данных: {e}")
176
  return False
177
 
178
  def get_page_text(self, page_num: int) -> str:
@@ -234,42 +227,37 @@ class VectorRAGSystem:
234
  return []
235
 
236
  def vector_search(self, query: str, k: int = 20) -> List[Tuple[Dict, float]]:
237
- """Векторный поиск по запросу"""
238
- if not self.faiss_index or not self.client:
239
- print("⚠️ FAISS индекс или OpenAI клиент недоступны")
240
- return []
241
-
242
  try:
243
- # Создаем эмбеддинг для запроса
244
  response = self.client.embeddings.create(
245
  model=self.embedding_model,
246
  input=[query]
247
  )
 
 
 
248
 
249
- query_embedding = np.array(response.data[0].embedding, dtype=np.float32)
250
- query_embedding = query_embedding.reshape(1, -1)
251
-
252
- # Нормализуем для Inner Product
253
- faiss.normalize_L2(query_embedding)
254
-
255
- # Поиск в FAISS индексе
256
- scores, indices = self.faiss_index.search(query_embedding, k)
257
-
258
- # Формируем результаты с parent-page enrichment
259
  results = []
260
  for score, idx in zip(scores[0], indices[0]):
261
  if 0 <= idx < len(self.chunks):
262
- chunk = self.chunks[idx].copy()
263
- # Получаем полный текст страницы для parent-page enrichment
264
- page_text = self.get_page_text(chunk["page"])
265
- chunk["text"] = page_text if page_text else chunk["text"]
266
- results.append((chunk, float(score)))
267
-
 
 
 
 
 
 
 
 
268
  return results
269
-
270
  except Exception as e:
271
- print(f"❌ Ошибка векторного поиска: {e}")
272
- print("⚠️ Переход на поиск без векторов невозможен")
273
  return []
274
 
275
  def rerank_with_llm(self, query: str, chunks: List[Tuple[Dict, float]]) -> List[Tuple[Dict, float]]:
 
119
  return False
120
 
121
  def load_vector_data(self) -> bool:
122
+ """Загрузка векторных данных и сохранение полной metadata_list с caption."""
123
  try:
124
  print("🔄 Попытка загрузки векторных данных...")
125
 
126
+ faiss_file = "chunks_flatip.faiss"
127
+ metadata_file = "metadata.json"
 
 
128
  if not all(os.path.exists(f) for f in [faiss_file, metadata_file]):
129
+ print("📁 Векторные файлы не найдены")
130
  return False
131
 
132
+ # 1) Читаем весь список метаданных, сохраняем его
133
  with open(metadata_file, 'r', encoding='utf-8') as f:
134
  metadata_list = json.load(f)
135
+ self.metadata_list = metadata_list
136
 
137
+ # 2) Строим self.chunks, сохраняя каждый item целиком
138
  self.chunks = []
139
  for i, item in enumerate(metadata_list):
 
140
  chunk_id = item.get("chunk_id",
141
  item.get("table_id",
142
  item.get("img_id", None)))
 
143
  self.chunks.append({
144
  "page": item["page"],
145
  "chunk_id": chunk_id,
146
  "chunk_index": i,
147
+ "text": "", # заполним в vector_search
148
+ "metadata": item # здесь есть caption, type и т.д.
149
  })
150
 
151
+ # 3) Загружаем FAISS-индекс
 
 
 
152
  if HAS_FAISS:
153
  self.faiss_index = faiss.read_index(faiss_file)
154
 
155
+ # 4) Загружаем PDF для parent-page enrichment
156
  pdf_path = "data/Сбер 2023.pdf"
157
  if os.path.exists(pdf_path):
158
+ import fitz
159
  self.pdf_doc = fitz.open(pdf_path)
160
  print(f"✅ PDF загружен: {self.pdf_doc.page_count} страниц")
161
  else:
162
+ print("❌ PDF не найден для enrichment")
163
  self.pdf_doc = None
164
 
165
+ print(f"✅ Загружены векторы: {len(self.chunks)} чанков")
166
  return True
 
167
  except Exception as e:
168
+ print(f"❌ Ошибка load_vector_data: {e}")
169
  return False
170
 
171
  def get_page_text(self, page_num: int) -> str:
 
227
  return []
228
 
229
  def vector_search(self, query: str, k: int = 20) -> List[Tuple[Dict, float]]:
230
+ """Векторный поиск + enrichment с caption из metadata_list."""
 
 
 
 
231
  try:
 
232
  response = self.client.embeddings.create(
233
  model=self.embedding_model,
234
  input=[query]
235
  )
236
+ q_emb = np.array(response.data[0].embedding, dtype=np.float32).reshape(1, -1)
237
+ faiss.normalize_L2(q_emb)
238
+ scores, indices = self.faiss_index.search(q_emb, k)
239
 
 
 
 
 
 
 
 
 
 
 
240
  results = []
241
  for score, idx in zip(scores[0], indices[0]):
242
  if 0 <= idx < len(self.chunks):
243
+ record = self.chunks[idx].copy()
244
+ meta_item = self.metadata_list[idx]
245
+
246
+ # базовый текст страницы
247
+ page_text = self.get_page_text(record["page"]) or ""
248
+
249
+ # если это картинка и есть caption — добавляем его сверху
250
+ if meta_item.get("type") == "image" and meta_item.get("caption"):
251
+ caption = meta_item["caption"]
252
+ record["text"] = caption + "\n\n" + page_text
253
+ else:
254
+ record["text"] = page_text
255
+
256
+ results.append((record, float(score)))
257
  return results
258
+
259
  except Exception as e:
260
+ print(f"❌ Ошибка vector_search: {e}")
 
261
  return []
262
 
263
  def rerank_with_llm(self, query: str, chunks: List[Tuple[Dict, float]]) -> List[Tuple[Dict, float]]: