MrSimple07 commited on
Commit
f3e59e1
·
1 Parent(s): 57e4dbd

simplest version

Browse files
Files changed (3) hide show
  1. documents_prep.py +0 -1
  2. index_retriever.py +3 -3
  3. utils.py +113 -155
documents_prep.py CHANGED
@@ -20,7 +20,6 @@ def chunk_text_documents(documents):
20
  chunked = []
21
 
22
  for doc in documents:
23
- # Skip duplicates or too-short content
24
  text_normalized = doc.text.strip()
25
  if len(text_normalized) < 50 or text_normalized in seen_texts:
26
  continue
 
20
  chunked = []
21
 
22
  for doc in documents:
 
23
  text_normalized = doc.text.strip()
24
  if len(text_normalized) < 50 or text_normalized in seen_texts:
25
  continue
index_retriever.py CHANGED
@@ -35,19 +35,19 @@ def create_query_engine(vector_index):
35
  # Vector retriever
36
  vector_retriever = VectorIndexRetriever(
37
  index=vector_index,
38
- similarity_top_k=30
39
  )
40
 
41
  # BM25 retriever
42
  bm25_retriever = BM25Retriever.from_defaults(
43
  docstore=vector_index.docstore,
44
- similarity_top_k=30
45
  )
46
 
47
  # Hybrid fusion
48
  hybrid_retriever = QueryFusionRetriever(
49
  [vector_retriever, bm25_retriever],
50
- similarity_top_k=40,
51
  num_queries=1
52
  )
53
 
 
35
  # Vector retriever
36
  vector_retriever = VectorIndexRetriever(
37
  index=vector_index,
38
+ similarity_top_k=40
39
  )
40
 
41
  # BM25 retriever
42
  bm25_retriever = BM25Retriever.from_defaults(
43
  docstore=vector_index.docstore,
44
+ similarity_top_k=40
45
  )
46
 
47
  # Hybrid fusion
48
  hybrid_retriever = QueryFusionRetriever(
49
  [vector_retriever, bm25_retriever],
50
+ similarity_top_k=50,
51
  num_queries=1
52
  )
53
 
utils.py CHANGED
@@ -40,74 +40,96 @@ def preprocess_query(question):
40
  import re
41
 
42
  question_lower = question.lower()
43
-
44
- # Extract document ID and normalize
45
- doc_match = re.search(r'(гост|нп|му)\s*р?\s*[№-]*\s*([0-9\.-]+)', question_lower)
46
-
47
  enhanced_query = question
48
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
49
  if doc_match:
50
- doc_type = doc_match.group(1).upper()
51
- doc_num = doc_match.group(2)
52
-
53
- # Add normalized versions
54
- enhanced_query += f" {doc_type} Р {doc_num}"
55
 
56
  return enhanced_query
57
 
58
  def answer_question(question, query_engine, reranker):
59
  try:
60
- log_message(f"Query: {question}")
 
61
 
62
  enhanced_query = preprocess_query(question)
63
- if enhanced_query != question:
64
- log_message(f"Enhanced query: {enhanced_query}")
 
 
 
65
 
66
  retrieved = query_engine.retriever.retrieve(enhanced_query)
67
- log_message(f"Retrieved {len(retrieved)} nodes")
68
 
 
69
  doc_stats = {}
70
  for n in retrieved:
71
  doc_id = n.metadata.get('document_id', 'unknown')
72
- doc_type = n.metadata.get('type', 'text')
73
 
74
- if doc_id not in doc_stats:
75
- doc_stats[doc_id] = {'tables': set(), 'text': 0, 'images': 0}
76
 
77
- if doc_type == 'table':
78
  table_id = n.metadata.get('table_identifier', n.metadata.get('table_number', '?'))
79
- doc_stats[doc_id]['tables'].add(table_id)
80
- elif doc_type == 'image':
81
- doc_stats[doc_id]['images'] += 1
82
  else:
83
- doc_stats[doc_id]['text'] += 1
84
 
85
  for doc_id in sorted(doc_stats.keys()):
86
  stats = doc_stats[doc_id]
87
- parts = []
88
  if stats['tables']:
89
- parts.append(f"tables={list(stats['tables'])[:5]}")
90
- if stats['text']:
91
- parts.append(f"text={stats['text']}")
92
- if stats['images']:
93
- parts.append(f"images={stats['images']}")
94
- log_message(f" {doc_id}: {', '.join(parts)}")
95
-
96
- doc_ids = [n.metadata.get('document_id', 'unknown') for n in retrieved]
97
- table_nums = [n.metadata.get('table_number', '') for n in retrieved if n.metadata.get('type') == 'table']
98
- log_message(f"Retrieved from documents: {set(doc_ids)}")
99
- if table_nums:
100
- log_message(f"Retrieved tables: {set(table_nums)}")
101
 
102
- reranked = rerank_nodes(question, retrieved, reranker, top_k=25)
103
- log_message(f"Reranked to {len(reranked)} nodes")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
104
 
105
- doc_ids_reranked = [n.metadata.get('document_id', 'unknown') for n in reranked]
106
- table_nums_reranked = [n.metadata.get('table_number', '') for n in reranked if n.metadata.get('type') == 'table']
107
- log_message(f"After reranking - documents: {set(doc_ids_reranked)}")
108
- if table_nums_reranked:
109
- log_message(f"After reranking - tables: {set(table_nums_reranked)}")
110
 
 
111
  context_parts = []
112
  for n in reranked:
113
  meta = n.metadata
@@ -115,103 +137,48 @@ def answer_question(question, query_engine, reranker):
115
  doc_type = meta.get('type', 'text')
116
 
117
  if doc_type == 'table':
118
- table_num = meta.get('table_number', 'unknown')
119
  title = meta.get('table_title', '')
120
- source_label = f"[ТАБЛИЦА {table_num} - {doc_id}]"
121
  if title:
122
  source_label += f" {title}"
123
- elif doc_type == 'image':
124
- img_num = meta.get('image_number', 'unknown')
125
- source_label = f"[РИСУНОК {img_num} - {doc_id}]"
126
  else:
127
- section = meta.get('section_id', '')
128
- source_label = f"[{doc_id} - {section}]"
129
 
130
- context_parts.append(f"{source_label}\n{n.text}")
131
 
132
- context = "\n\n" + ("="*70 + "\n\n").join(context_parts)
133
- from config import CUSTOM_PROMPT
134
- prompt = f"""
135
- Вы являетесь высокоспециализированным Ассистентом для анализа нормативных документов (AIEXP). Ваша цель - предоставлять точные, корректные и контекстно релевантные ответы исключительно на основе предоставленного контекста из нормативной документации.
136
-
137
- ПРАВИЛА АНАЛИЗА ЗАПРОСА:
138
-
139
- 1. ПРЯМЫЕ ВОПРОСЫ БЕЗ ДОКУМЕНТАЛЬНОГО КОНТЕКСТА:
140
- Если пользователь задает вопрос типа "В каких случаях могут быть признаны протоколы испытаний?" без предоставления дополнительных документов, найдите соответствующую информацию в доступном контексте и предоставьте полный ответ с указанием источников.
141
-
142
- 2. ОПРЕДЕЛЕНИЕ ТИПА ЗАДАЧИ:
143
-
144
- а) ПОИСК И ОТВЕТ НА ВОПРОС (ключевые слова: "в каких случаях", "когда", "кто", "что", "как", "почему"):
145
- - Найдите релевантную информацию в контексте
146
- - Предоставьте развернутый ответ
147
- - Обязательно укажите конкретные документы и разделы
148
- - Процитируйте ключевые положения
149
-
150
- б) КРАТКОЕ САММАРИ (ключевые слова: "кратко", "суммировать", "резюме", "основные моменты"):
151
- - Предоставьте структурированное резюме
152
- - Выделите ключевые требования
153
- - Используйте нумерованный список
154
-
155
- в) ПОИСК ДОКУМЕНТА И ПУНКТА (ключевые слова: "найти", "где", "какой документ", "в каком разделе"):
156
- - Укажите конкретный документ и структурное расположение
157
- - Предоставьте точные номера разделов/пунктов
158
-
159
- г) ПРОВЕРКА КОРРЕКТНОСТИ (ключевые слова: "правильно ли", "соответствует ли", "проверить"):
160
- - Четко укажите: "СООТВЕТСТВУЕТ" или "НЕ СООТВЕТСТВУЕТ"
161
- - Перечислите конкретные требования
162
-
163
- д) ПЛАН ДЕЙСТВИЙ (ключевые слова: "план", "алгоритм", "пошагово"):
164
- - Создайте пронумерованный план
165
- - Укажите ссылки на соответствующие пункты НД
166
-
167
- ПРАВИЛА ФОРМИРОВАНИЯ ОТВЕТОВ:
168
-
169
- Работай исключительно с информацией из предоставленного контекста. Запрещено использовать:
170
- - Общие знания
171
- - Информацию из интернета
172
- - Данные из предыдущих диалогов
173
- - Собственные предположения
174
-
175
- 1. СТРУКТУРА ОТВЕТА:
176
- - Начинайте с прямого ответа на вопрос
177
- - Затем указывайте нормативные основания
178
- - Завершайте ссылками на конкретные документы и разделы
179
-
180
- 2. РАБОТА С КОНТЕКСТОМ:
181
- - Если информация найдена в контексте - предоставьте полный ответ
182
- - Если информация не найдена: "Информация по вашему запросу не найдена в доступной нормативной документации"
183
- - Не делайте предположений за пределами контекста
184
- - Не используйте общие знания
185
-
186
- 3. ТЕРМИНОЛОГИЯ И ЦИТИРОВАНИЕ:
187
- - Сохраняйте официальную терминологию НД
188
- - Цитируйте точные формулировки ключевых требований
189
- - При множественных источниках - укажите все релевантные
190
-
191
- 4. ФОРМАТИРОВАНИЕ:
192
- - Для перечислений: используйте нумерованные списки
193
- - Выделяйте критически важные требования
194
- - Структурируйте ответ логически
195
-
196
- # КАК РАБОТАТЬ С ЗАПРОСОМ
197
 
198
- **Шаг 1:** Определи, что именно ищет пользователь (термин, требование, процедура, условие)
 
199
 
200
- **Шаг 2:** Найди релевантную информацию в контексте
201
 
202
- **Шаг 3:** Сформируй ответ:
203
- - Если нашел: укажи документ и пункт, процитируй нужную часть
204
- - Если не нашел: четко сообщи об отсутствии информации
 
205
 
206
- **Шаг 4:** При наличии нескольких источников:
207
- - Представь их последовательно с указанием источника каждого
208
- - Если источников много (>4) сначала дай их список, потом цитаты
209
 
210
  КОНТЕКСТ:
211
  {context}
212
 
213
  ВОПРОС: {question}
214
- """
 
 
 
 
 
 
 
215
 
216
  response = query_engine.query(prompt)
217
  sources = format_sources(reranked)
@@ -223,54 +190,45 @@ def answer_question(question, query_engine, reranker):
223
  import traceback
224
  log_message(traceback.format_exc())
225
  return f"Ошибка: {e}", ""
226
-
227
- def rerank_nodes(query, nodes, reranker, top_k=20, min_score=0.3):
228
- """Rerank nodes with diversity - MORE LENIENT"""
229
  if not nodes:
230
  return []
231
 
232
- # Score all nodes
233
  pairs = [[query, n.text] for n in nodes]
234
  scores = reranker.predict(pairs)
235
 
236
- # Sort by score
237
  scored = sorted(zip(nodes, scores), key=lambda x: x[1], reverse=True)
238
 
239
- # More lenient threshold
 
 
240
  filtered = [(n, s) for n, s in scored if s >= min_score]
241
 
242
  if not filtered:
243
- # Fallback: take top 50% if nothing passes threshold
244
- cutoff = max(scores) * 0.5
245
  filtered = [(n, s) for n, s in scored if s >= cutoff][:top_k]
246
 
247
- # Log top scores for debugging
248
- log_message(f"Top 5 reranking scores: {[f'{s:.3f}' for _, s in scored[:5]]}")
 
 
 
 
 
249
 
250
- # Diversity selection - but prioritize tables if query mentions them
251
  selected = []
252
- seen_docs = set()
253
- table_nodes = []
254
- other_nodes = []
255
-
256
- for node, score in filtered:
257
- if node.metadata.get('type') == 'table':
258
- table_nodes.append((node, score))
259
- else:
260
- other_nodes.append((node, score))
261
 
262
- # If query mentions "таблица", prioritize table nodes
263
- if 'таблиц' in query.lower():
264
- combined = table_nodes + other_nodes
265
- else:
266
- combined = filtered
267
 
268
- for node, score in combined[:top_k]:
269
- if len(selected) >= top_k:
270
- break
271
- selected.append(node)
272
- seen_docs.add(node.metadata.get('document_id', 'unknown'))
273
 
274
- log_message(f"Reranked: {len(filtered)} → {len(selected)} (from {len(seen_docs)} docs)")
275
 
276
- return selected
 
40
  import re
41
 
42
  question_lower = question.lower()
 
 
 
 
43
  enhanced_query = question
44
 
45
+ # Detect "list all tables" queries - handle differently
46
+ if any(phrase in question_lower for phrase in ['какие таблиц', 'список таблиц', 'перечисл', 'все таблиц']):
47
+ # For listing queries, just extract document ID
48
+ doc_match = re.search(r'(гост|нп|му)[^\s]*\s*р?\s*[№-]*\s*([0-9\.-]+)', question_lower)
49
+ if doc_match:
50
+ doc_id = f"{doc_match.group(1).upper()} Р {doc_match.group(2)}"
51
+ enhanced_query = f"документ {doc_id} таблица"
52
+ return enhanced_query
53
+
54
+ # For specific table queries
55
+ table_match = re.search(r'табли[цу]\w*\s+(?:№|номер)?\s*([а-яa-z0-9\.]+)', question_lower)
56
+ if table_match:
57
+ table_num = table_match.group(1).upper()
58
+ enhanced_query += f" таблица {table_num}"
59
+
60
+ # Document detection
61
+ doc_match = re.search(r'(гост|нп|му)[^\s]*\s*р?\s*[№-]*\s*([0-9\.-]+)', question_lower)
62
  if doc_match:
63
+ doc_id = f"{doc_match.group(1).upper()} Р {doc_match.group(2)}"
64
+ enhanced_query += f" документ {doc_id}"
 
 
 
65
 
66
  return enhanced_query
67
 
68
  def answer_question(question, query_engine, reranker):
69
  try:
70
+ log_message(f"\n{'='*70}")
71
+ log_message(f"QUERY: {question}")
72
 
73
  enhanced_query = preprocess_query(question)
74
+ log_message(f"Enhanced: {enhanced_query}")
75
+
76
+ # Detect listing queries - need MORE chunks
77
+ is_listing_query = any(phrase in question.lower()
78
+ for phrase in ['какие таблиц', 'список', 'перечисл', 'все таблиц'])
79
 
80
  retrieved = query_engine.retriever.retrieve(enhanced_query)
81
+ log_message(f"\nRETRIEVED: {len(retrieved)} nodes")
82
 
83
+ # Log retrieved docs
84
  doc_stats = {}
85
  for n in retrieved:
86
  doc_id = n.metadata.get('document_id', 'unknown')
87
+ doc_group = n.metadata.get('document_group', doc_id)
88
 
89
+ if doc_group not in doc_stats:
90
+ doc_stats[doc_group] = {'tables': set(), 'text': 0}
91
 
92
+ if n.metadata.get('type') == 'table':
93
  table_id = n.metadata.get('table_identifier', n.metadata.get('table_number', '?'))
94
+ doc_stats[doc_group]['tables'].add(table_id)
 
 
95
  else:
96
+ doc_stats[doc_group]['text'] += 1
97
 
98
  for doc_id in sorted(doc_stats.keys()):
99
  stats = doc_stats[doc_id]
100
+ log_message(f" {doc_id}: {len(stats['tables'])} tables, {stats['text']} text")
101
  if stats['tables']:
102
+ log_message(f" Tables: {sorted(stats['tables'])}")
 
 
 
 
 
 
 
 
 
 
 
103
 
104
+ # Adjust reranking based on query type
105
+ if is_listing_query:
106
+ reranked = rerank_nodes(question, retrieved, reranker, top_k=50, min_score=0.2)
107
+ else:
108
+ reranked = rerank_nodes(question, retrieved, reranker, top_k=25, min_score=0.3)
109
+
110
+ log_message(f"\nRERANKED: {len(reranked)} nodes")
111
+
112
+ # Log reranked
113
+ doc_stats_reranked = {}
114
+ for n in reranked:
115
+ doc_group = n.metadata.get('document_group', n.metadata.get('document_id', 'unknown'))
116
+
117
+ if doc_group not in doc_stats_reranked:
118
+ doc_stats_reranked[doc_group] = {'tables': set(), 'text': 0}
119
+
120
+ if n.metadata.get('type') == 'table':
121
+ table_id = n.metadata.get('table_identifier', n.metadata.get('table_number', '?'))
122
+ doc_stats_reranked[doc_group]['tables'].add(table_id)
123
+ else:
124
+ doc_stats_reranked[doc_group]['text'] += 1
125
 
126
+ for doc_id in sorted(doc_stats_reranked.keys()):
127
+ stats = doc_stats_reranked[doc_id]
128
+ log_message(f" {doc_id}: {len(stats['tables'])} tables, {stats['text']} text")
129
+ if stats['tables']:
130
+ log_message(f" Tables: {sorted(stats['tables'])}")
131
 
132
+ # Build context
133
  context_parts = []
134
  for n in reranked:
135
  meta = n.metadata
 
137
  doc_type = meta.get('type', 'text')
138
 
139
  if doc_type == 'table':
140
+ table_id = meta.get('table_identifier', meta.get('table_number', 'unknown'))
141
  title = meta.get('table_title', '')
142
+ source_label = f"[{doc_id} - Таблица {table_id}]"
143
  if title:
144
  source_label += f" {title}"
 
 
 
145
  else:
146
+ source_label = f"[{doc_id}]"
 
147
 
148
+ context_parts.append(f"{source_label}\n{n.text[:500]}") # Limit context per chunk
149
 
150
+ context = "\n\n" + ("="*50 + "\n\n").join(context_parts)
151
+
152
+ # Adjust prompt for listing queries
153
+ if is_listing_query:
154
+ prompt = f"""Контекст содержит информацию о таблицах из документов.
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
155
 
156
+ КОНТЕКСТ:
157
+ {context}
158
 
159
+ ВОПРОС: {question}
160
 
161
+ ИНСТРУКЦИИ:
162
+ 1. Перечисли ВСЕ таблицы, найденные в контексте для запрошенного документа
163
+ 2. Укажи номер таблицы и название (если есть)
164
+ 3. Если таблиц нет - скажи прямо
165
 
166
+ ОТВЕТ (список таблиц):"""
167
+ else:
168
+ prompt = f"""Ты эксперт по технической документации.
169
 
170
  КОНТЕКСТ:
171
  {context}
172
 
173
  ВОПРОС: {question}
174
+
175
+ ИНСТРУКЦИИ:
176
+ 1. Отвечай ТОЛЬКО на основе контекста
177
+ 2. Укажи источник (документ, таблицу)
178
+ 3. Если нужно показать содержимое таблицы - покажи ВСЕ данные
179
+ 4. Если информации нет - скажи прямо
180
+
181
+ ОТВЕТ:"""
182
 
183
  response = query_engine.query(prompt)
184
  sources = format_sources(reranked)
 
190
  import traceback
191
  log_message(traceback.format_exc())
192
  return f"Ошибка: {e}", ""
193
+ def rerank_nodes(query, nodes, reranker, top_k=25, min_score=0.3):
194
+ """Rerank with document grouping awareness"""
 
195
  if not nodes:
196
  return []
197
 
 
198
  pairs = [[query, n.text] for n in nodes]
199
  scores = reranker.predict(pairs)
200
 
 
201
  scored = sorted(zip(nodes, scores), key=lambda x: x[1], reverse=True)
202
 
203
+ log_message(f"Top 10 reranking scores: {[f'{s:.3f}' for _, s in scored[:10]]}")
204
+
205
+ # More lenient filtering
206
  filtered = [(n, s) for n, s in scored if s >= min_score]
207
 
208
  if not filtered:
209
+ cutoff = max(scores) * 0.4
 
210
  filtered = [(n, s) for n, s in scored if s >= cutoff][:top_k]
211
 
212
+ # Group by document for diversity
213
+ doc_groups = {}
214
+ for node, score in filtered:
215
+ doc_group = node.metadata.get('document_group', node.metadata.get('document_id', 'unknown'))
216
+ if doc_group not in doc_groups:
217
+ doc_groups[doc_group] = []
218
+ doc_groups[doc_group].append((node, score))
219
 
220
+ # Take top chunks from each document group
221
  selected = []
222
+ group_limits = max(3, top_k // max(1, len(doc_groups)))
 
 
 
 
 
 
 
 
223
 
224
+ for doc_group in doc_groups:
225
+ selected.extend([n for n, s in doc_groups[doc_group][:group_limits]])
 
 
 
226
 
227
+ # Fill remaining slots with highest scores
228
+ if len(selected) < top_k:
229
+ remaining = [n for n, s in filtered if n not in selected]
230
+ selected.extend(remaining[:top_k - len(selected)])
 
231
 
232
+ log_message(f"Reranked: {len(filtered)} → {len(selected)} (from {len(doc_groups)} doc groups)")
233
 
234
+ return selected[:top_k]