0cf3297290
- Краулеры: smart_crawler.py, regional_crawler.py - Аудит: audit_orel_to_excel.py, audit_chukotka_to_excel.py - РКН проверка: check_rkn_registry.py, recheck_unclear_rkn.py - Отчёты: create_orel_horizontal_report.py - Обработка: process_all_hotels_embeddings.py - Документация: README.md, DB_SCHEMA_REFERENCE.md
397 lines
17 KiB
Python
397 lines
17 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Скрипт для обработки всех отелей:
|
||
1. Создание chunks из hotel_website_processed
|
||
2. Генерация эмбеддингов через BGE-M3 API
|
||
3. Сохранение в hotel_website_chunks с metadata
|
||
"""
|
||
|
||
import psycopg2
|
||
from urllib.parse import unquote
|
||
import requests
|
||
import json
|
||
import time
|
||
import logging
|
||
from typing import List, Dict, Tuple
|
||
import uuid
|
||
|
||
# Настройка логирования
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format='%(asctime)s - %(levelname)s - %(message)s',
|
||
handlers=[
|
||
logging.FileHandler('embeddings_processing.log'),
|
||
logging.StreamHandler()
|
||
]
|
||
)
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# Конфигурация
|
||
BGE_API_URL = "http://147.45.146.17:8002/embed"
|
||
BGE_API_KEY = "22564b177aa73b6ac0b8642d7773350ff4c01d4983f028beff15ea247f09fa89"
|
||
CHUNK_SIZE = 600
|
||
CHUNK_OVERLAP = 100
|
||
BATCH_SIZE = 8 # Размер батча для API (уменьшен из-за перегрузки)
|
||
MAX_RETRIES = 3 # Количество попыток при ошибке
|
||
|
||
class EmbeddingProcessor:
|
||
def __init__(self):
|
||
self.conn = None
|
||
self.cur = None
|
||
self.connect_db()
|
||
|
||
def connect_db(self):
|
||
"""Подключение к базе данных"""
|
||
try:
|
||
self.conn = psycopg2.connect(
|
||
host='147.45.189.234',
|
||
port=5432,
|
||
database='default_db',
|
||
user='gen_user',
|
||
password=unquote('2~~9_%5EkVsU%3F2%5CS')
|
||
)
|
||
self.conn.autocommit = True
|
||
self.cur = self.conn.cursor()
|
||
logger.info("✅ Подключение к БД установлено")
|
||
except Exception as e:
|
||
logger.error(f"❌ Ошибка подключения к БД: {e}")
|
||
raise
|
||
|
||
def get_hotel_info(self, hotel_id: str) -> Dict:
|
||
"""Получение информации об отеле из hotel_main"""
|
||
try:
|
||
self.cur.execute("""
|
||
SELECT id, full_name, region_name
|
||
FROM hotel_main
|
||
WHERE id = %s;
|
||
""", (hotel_id,))
|
||
|
||
result = self.cur.fetchone()
|
||
if result:
|
||
return {
|
||
'hotel_id': result[0],
|
||
'hotel_name': result[1],
|
||
'region_name': result[2]
|
||
}
|
||
return None
|
||
except Exception as e:
|
||
logger.error(f"❌ Ошибка получения информации об отеле {hotel_id}: {e}")
|
||
return None
|
||
|
||
def create_chunks_from_text(self, text: str, hotel_id: str, url: str, raw_page_id: int) -> List[Dict]:
|
||
"""Создание chunks из текста"""
|
||
if not text or len(text.strip()) < 50:
|
||
return []
|
||
|
||
chunks = []
|
||
start = 0
|
||
|
||
while start < len(text):
|
||
end = start + CHUNK_SIZE
|
||
chunk_text = text[start:end]
|
||
|
||
if end < len(text):
|
||
# Ищем хорошую точку разрыва (конец предложения)
|
||
last_period = chunk_text.rfind('.')
|
||
last_newline = chunk_text.rfind('\n')
|
||
break_point = max(last_period, last_newline)
|
||
|
||
if break_point > start + CHUNK_SIZE // 2:
|
||
chunk_text = text[start:start + break_point + 1]
|
||
end = start + break_point + 1
|
||
|
||
# Создаём metadata для chunk
|
||
chunk_metadata = {
|
||
'hotel_id': hotel_id,
|
||
'url': url,
|
||
'raw_page_id': raw_page_id,
|
||
'chunk_start': start,
|
||
'chunk_end': end,
|
||
'chunk_length': len(chunk_text)
|
||
}
|
||
|
||
chunks.append({
|
||
'id': str(uuid.uuid4()),
|
||
'text': chunk_text.strip(),
|
||
'metadata': chunk_metadata
|
||
})
|
||
|
||
# Следующий chunk с перекрытием
|
||
start = end - CHUNK_OVERLAP
|
||
if start >= len(text):
|
||
break
|
||
|
||
return chunks
|
||
|
||
def generate_embeddings_batch(self, texts: List[str]) -> List[List[float]]:
|
||
"""Генерация эмбеддингов батчем через API с retry логикой"""
|
||
for attempt in range(MAX_RETRIES):
|
||
try:
|
||
headers = {
|
||
"X-API-Key": BGE_API_KEY,
|
||
"Content-Type": "application/json"
|
||
}
|
||
|
||
payload = {
|
||
"text": texts
|
||
}
|
||
|
||
# Увеличиваем таймаут для больших батчей
|
||
timeout = 120 if len(texts) > 20 else 60
|
||
|
||
response = requests.post(BGE_API_URL, json=payload, headers=headers, timeout=timeout)
|
||
|
||
if response.status_code == 200:
|
||
result = response.json()
|
||
embeddings = result.get('embeddings', [])
|
||
if len(embeddings) == len(texts):
|
||
return embeddings
|
||
else:
|
||
logger.warning(f"⚠️ Неполный ответ API: {len(embeddings)}/{len(texts)} эмбеддингов")
|
||
if attempt < MAX_RETRIES - 1:
|
||
logger.info(f"🔄 Повторная попытка {attempt + 2}/{MAX_RETRIES}")
|
||
time.sleep(5) # Пауза перед повтором
|
||
continue
|
||
else:
|
||
logger.error(f"❌ Ошибка API: {response.status_code} - {response.text}")
|
||
if attempt < MAX_RETRIES - 1:
|
||
logger.info(f"🔄 Повторная попытка {attempt + 2}/{MAX_RETRIES}")
|
||
time.sleep(10) # Пауза перед повтором
|
||
continue
|
||
|
||
except requests.exceptions.Timeout:
|
||
logger.warning(f"⚠️ Таймаут API (попытка {attempt + 1}/{MAX_RETRIES})")
|
||
if attempt < MAX_RETRIES - 1:
|
||
logger.info(f"🔄 Повторная попытка {attempt + 2}/{MAX_RETRIES}")
|
||
time.sleep(10) # Пауза перед повтором
|
||
continue
|
||
except Exception as e:
|
||
logger.error(f"❌ Ошибка генерации эмбеддингов (попытка {attempt + 1}/{MAX_RETRIES}): {e}")
|
||
if attempt < MAX_RETRIES - 1:
|
||
logger.info(f"🔄 Повторная попытка {attempt + 2}/{MAX_RETRIES}")
|
||
time.sleep(5) # Пауза перед повтором
|
||
continue
|
||
|
||
logger.error(f"❌ Не удалось получить эмбеддинги после {MAX_RETRIES} попыток")
|
||
return []
|
||
|
||
def save_chunks_to_db(self, chunks: List[Dict], hotel_info: Dict):
|
||
"""Сохранение chunks в базу данных"""
|
||
try:
|
||
# Разбиваем chunks на батчи для API
|
||
all_embeddings = []
|
||
|
||
for i in range(0, len(chunks), BATCH_SIZE):
|
||
batch_chunks = chunks[i:i + BATCH_SIZE]
|
||
batch_texts = [chunk['text'] for chunk in batch_chunks]
|
||
|
||
logger.info(f" 🔄 Обрабатываем батч {i//BATCH_SIZE + 1}: {len(batch_texts)} chunks")
|
||
|
||
# Генерируем эмбеддинги для батча
|
||
batch_embeddings = self.generate_embeddings_batch(batch_texts)
|
||
|
||
if len(batch_embeddings) == len(batch_texts):
|
||
all_embeddings.extend(batch_embeddings)
|
||
logger.info(f" ✅ Батч успешно обработан")
|
||
else:
|
||
logger.error(f" ❌ Ошибка в батче: {len(batch_embeddings)}/{len(batch_texts)} эмбеддингов")
|
||
return False
|
||
|
||
# Небольшая пауза между батчами
|
||
if i + BATCH_SIZE < len(chunks):
|
||
time.sleep(1)
|
||
|
||
if len(all_embeddings) != len(chunks):
|
||
logger.error(f"❌ Количество эмбеддингов ({len(all_embeddings)}) не совпадает с количеством chunks ({len(chunks)})")
|
||
return False
|
||
|
||
# Обновляем metadata с информацией об отеле и сохраняем в БД
|
||
for i, chunk in enumerate(chunks):
|
||
chunk['metadata']['hotel_name'] = hotel_info['hotel_name']
|
||
chunk['metadata']['region_name'] = hotel_info['region_name']
|
||
|
||
# Сохраняем в БД
|
||
embedding_str = json.dumps(all_embeddings[i])
|
||
|
||
self.cur.execute("""
|
||
INSERT INTO hotel_website_chunks (id, text, metadata, embedding)
|
||
VALUES (%s, %s, %s, %s::vector)
|
||
ON CONFLICT (id) DO UPDATE SET
|
||
text = EXCLUDED.text,
|
||
metadata = EXCLUDED.metadata,
|
||
embedding = EXCLUDED.embedding;
|
||
""", (
|
||
chunk['id'],
|
||
chunk['text'],
|
||
json.dumps(chunk['metadata']),
|
||
embedding_str
|
||
))
|
||
|
||
logger.info(f"✅ Сохранено {len(chunks)} chunks для отеля {hotel_info['hotel_name'][:50]}...")
|
||
return True
|
||
|
||
except Exception as e:
|
||
logger.error(f"❌ Ошибка сохранения chunks: {e}")
|
||
return False
|
||
|
||
def process_hotel(self, hotel_id: str) -> bool:
|
||
"""Обработка одного отеля"""
|
||
try:
|
||
# Получаем информацию об отеле
|
||
hotel_info = self.get_hotel_info(hotel_id)
|
||
if not hotel_info:
|
||
logger.warning(f"⚠️ Отель {hotel_id} не найден в hotel_main")
|
||
return False
|
||
|
||
# Получаем страницы отеля
|
||
self.cur.execute("""
|
||
SELECT id, url, cleaned_text
|
||
FROM hotel_website_processed
|
||
WHERE hotel_id = %s
|
||
AND cleaned_text IS NOT NULL
|
||
AND LENGTH(cleaned_text) > 50
|
||
ORDER BY id;
|
||
""", (hotel_id,))
|
||
|
||
pages = self.cur.fetchall()
|
||
logger.info(f"🏨 Обрабатываем отель: {hotel_info['hotel_name'][:50]}...")
|
||
logger.info(f" 📄 Найдено {len(pages)} страниц")
|
||
|
||
total_chunks = 0
|
||
|
||
for page_id, url, text in pages:
|
||
# Создаём chunks
|
||
chunks = self.create_chunks_from_text(text, hotel_id, url, page_id)
|
||
|
||
if chunks:
|
||
# Сохраняем chunks
|
||
if self.save_chunks_to_db(chunks, hotel_info):
|
||
total_chunks += len(chunks)
|
||
logger.info(f" ✅ Страница {page_id}: {len(chunks)} chunks")
|
||
else:
|
||
logger.error(f" ❌ Ошибка сохранения chunks для страницы {page_id}")
|
||
|
||
logger.info(f"🎉 Отель {hotel_info['hotel_name'][:50]}... обработан: {total_chunks} chunks")
|
||
return True
|
||
|
||
except Exception as e:
|
||
logger.error(f"❌ Ошибка обработки отеля {hotel_id}: {e}")
|
||
return False
|
||
|
||
def get_hotels_to_process(self) -> List[str]:
|
||
"""Получение списка отелей для обработки"""
|
||
try:
|
||
# Получаем отели, которые есть в hotel_website_processed, но нет в chunks
|
||
self.cur.execute("""
|
||
SELECT DISTINCT p.hotel_id
|
||
FROM hotel_website_processed p
|
||
LEFT JOIN hotel_website_chunks c ON p.hotel_id::text = c.metadata->>'hotel_id'
|
||
WHERE p.cleaned_text IS NOT NULL
|
||
AND LENGTH(p.cleaned_text) > 50
|
||
AND c.id IS NULL
|
||
ORDER BY p.hotel_id;
|
||
""")
|
||
|
||
hotels = [row[0] for row in self.cur.fetchall()]
|
||
logger.info(f"📊 Найдено {len(hotels)} отелей для обработки")
|
||
return hotels
|
||
|
||
except Exception as e:
|
||
logger.error(f"❌ Ошибка получения списка отелей: {e}")
|
||
return []
|
||
|
||
def get_processing_stats(self) -> Dict:
|
||
"""Получение статистики обработки"""
|
||
try:
|
||
self.cur.execute("""
|
||
SELECT
|
||
COUNT(DISTINCT p.hotel_id) as total_hotels,
|
||
COUNT(p.id) as total_pages,
|
||
COUNT(DISTINCT c.metadata->>'hotel_id') as processed_hotels,
|
||
COUNT(c.id) as total_chunks
|
||
FROM hotel_website_processed p
|
||
LEFT JOIN hotel_website_chunks c ON p.hotel_id::text = c.metadata->>'hotel_id'
|
||
WHERE p.cleaned_text IS NOT NULL AND LENGTH(p.cleaned_text) > 50;
|
||
""")
|
||
|
||
result = self.cur.fetchone()
|
||
return {
|
||
'total_hotels': result[0],
|
||
'total_pages': result[1],
|
||
'processed_hotels': result[2],
|
||
'total_chunks': result[3]
|
||
}
|
||
|
||
except Exception as e:
|
||
logger.error(f"❌ Ошибка получения статистики: {e}")
|
||
return {}
|
||
|
||
def close(self):
|
||
"""Закрытие соединения с БД"""
|
||
if self.cur:
|
||
self.cur.close()
|
||
if self.conn:
|
||
self.conn.close()
|
||
|
||
def main():
|
||
"""Основная функция"""
|
||
logger.info("🚀 Запуск обработки отелей для эмбеддингов")
|
||
|
||
processor = EmbeddingProcessor()
|
||
|
||
try:
|
||
# Получаем статистику
|
||
stats = processor.get_processing_stats()
|
||
logger.info(f"📊 Текущая статистика:")
|
||
logger.info(f" Всего отелей: {stats.get('total_hotels', 0)}")
|
||
logger.info(f" Всего страниц: {stats.get('total_pages', 0)}")
|
||
logger.info(f" Обработано отелей: {stats.get('processed_hotels', 0)}")
|
||
logger.info(f" Всего chunks: {stats.get('total_chunks', 0)}")
|
||
|
||
# Получаем отели для обработки
|
||
hotels_to_process = processor.get_hotels_to_process()
|
||
|
||
if not hotels_to_process:
|
||
logger.info("✅ Все отели уже обработаны!")
|
||
return
|
||
|
||
# Обрабатываем отели
|
||
successful = 0
|
||
failed = 0
|
||
|
||
for i, hotel_id in enumerate(hotels_to_process, 1):
|
||
logger.info(f"\n🔄 Обрабатываем отель {i}/{len(hotels_to_process)}: {hotel_id}")
|
||
|
||
start_time = time.time()
|
||
|
||
if processor.process_hotel(hotel_id):
|
||
successful += 1
|
||
processing_time = time.time() - start_time
|
||
logger.info(f"✅ Успешно за {processing_time:.2f} сек")
|
||
else:
|
||
failed += 1
|
||
logger.error(f"❌ Ошибка обработки")
|
||
|
||
# Показываем прогресс каждые 10 отелей
|
||
if i % 10 == 0:
|
||
logger.info(f"\n📈 Прогресс: {i}/{len(hotels_to_process)} отелей")
|
||
logger.info(f" ✅ Успешно: {successful}")
|
||
logger.info(f" ❌ Ошибок: {failed}")
|
||
|
||
# Финальная статистика
|
||
final_stats = processor.get_processing_stats()
|
||
logger.info(f"\n🎉 ОБРАБОТКА ЗАВЕРШЕНА!")
|
||
logger.info(f"📊 Финальная статистика:")
|
||
logger.info(f" Обработано отелей: {final_stats.get('processed_hotels', 0)}")
|
||
logger.info(f" Всего chunks: {final_stats.get('total_chunks', 0)}")
|
||
logger.info(f" ✅ Успешно: {successful}")
|
||
logger.info(f" ❌ Ошибок: {failed}")
|
||
|
||
except Exception as e:
|
||
logger.error(f"❌ Критическая ошибка: {e}")
|
||
finally:
|
||
processor.close()
|
||
|
||
if __name__ == "__main__":
|
||
main()
|