main.py - master - IvanMysin/Topics

main.py
187 строк9 KB
ivan
Working on preport scripts
18 фев 2026, 14:27
18 фев 2026, 14:275b40c9a
import pandas as pd
import numpy as np
from tqdm.auto import tqdm
import argparse
import sys
import os
from datetime import datetime

# Добавляем src в путь для импорта
sys.path.append(os.path.join(os.path.dirname(__file__), 'src'))

from embedding_manager import EmbeddingManager
from clustering.clustering_factory import ClusteringFactory
from report_generator import ReportGenerator
from results_manager import ResultsManager
from config import RESULTS_DIR, DB_CONFIG, LLM_CONFIG


def main():
    parser = argparse.ArgumentParser(description='Кластеризация научных статей')
    parser.add_argument('--force-update', action='store_true',
                        help='Принудительное обновление всех эмбеддингов')
    parser.add_argument('--clusters', type=int,
                        help='Количество кластеров (игнорируется для DBSCAN)')
    parser.add_argument('--max-documents', type=int,
                        help='Ограничение количества документов для тестирования')
    parser.add_argument('--method', type=str, default='xmeans',
                        choices=ClusteringFactory.get_available_methods(),
                        help=f'Метод кластеризации (по умолчанию: xmeans)')
    parser.add_argument('--create-dendrogram', action='store_true',
                        help='Создать дендрограмму для иерархической кластеризации')
    parser.add_argument('--create-visualization', action='store_true',
                        help='Создать визуализацию кластеров')
    parser.add_argument('--eps', type=float,
                        help='Значение eps для DBSCAN (если не задано, подбирается автоматически)')
    parser.add_argument('--min-samples', type=int,
                        help='min_samples для DBSCAN')
    parser.add_argument('--do-descriptions', action='store_false',
                        help='Отключить генерацию описаний кластеров с помощью LLM')
    parser.add_argument('--ollama-model', type=str,
                        help='Модель Ollama для генерации описаний')
    parser.add_argument('--output-format', type=str, default='both',
                        choices=['markdown', 'html', 'both'],
                        help='Формат вывода отчетов (markdown, html, both)')
    parser.add_argument('--keyword-method', type=str, default='tfidf',
                        choices=['tfidf', 'yake', 'keybert'],
                        help='Метод извлечения ключевых слов (tfidf, yake, keybert)')
    parser.add_argument('--date-from', type=str,
                        help='Фильтр: начальная дата публикации (формат: YYYY-MM-DD)')
    parser.add_argument('--date-to', type=str,
                        help='Фильтр: конечная дата публикации (формат: YYYY-MM-DD)')
    parser.add_argument('--journals-file', type=str,
                        help='Путь к файлу со списком журналов (по одному названию на строку)')

    args = parser.parse_args()

    try:
        # Инициализация менеджеров
        print("🔄 Инициализация менеджеров...")
        embedding_manager = EmbeddingManager()
        clustering_engine = ClusteringFactory.create_clusterer(args.method, args.keyword_method)
        results_manager = ResultsManager()
        report_generator = ReportGenerator(keyword_method=args.keyword_method)



        # Загрузка документов с фильтрами
        df = embedding_manager.load_documents_from_db(
            date_from=args.date_from,
            date_to=args.date_to,
            journals_file=args.journals_file
        )

        if args.max_documents:
            df = df.head(args.max_documents)
            print(f"🔬 Ограничение: используется {len(df)} документов")

        # Обновление эмбеддингов
        updated_count = embedding_manager.update_embeddings(df, args.force_update)

        if updated_count > 0:
            print(f"🔄 Обновлено {updated_count} эмбеддингов")

        # Получение всех эмбеддингов
        embeddings, valid_ids = embedding_manager.get_all_embeddings(df)

        if len(embeddings) == 0:
            print("❌ Не удалось загрузить эмбеддинги. Завершение работы.")
            return

        # Фильтруем DataFrame по valid_ids
        valid_df = df[df[DB_CONFIG['Articles']['id_column']].astype(str).isin(valid_ids)].copy()

        if len(valid_df) == 0:
            print("❌ Нет валидных документов для кластеризации.")
            return

        # Настройка параметров для DBSCAN
        if args.method == 'dbscan':
            if args.eps is not None:
                print(f"🎯 Используется заданное eps: {args.eps}")
                clustering_engine.optimal_eps = args.eps
            if args.min_samples is not None:
                print(f"🎯 Используется заданное min_samples: {args.min_samples}")
                clustering_engine.config["min_samples"] = args.min_samples

        # Кластеризация
        n_clusters = args.clusters
        labels, clustering_model = clustering_engine.fit_predict(embeddings, n_clusters)

        # Создание дендрограммы для иерархической кластеризации
        if args.method == 'hierarchical' and args.create_dendrogram:
            from clustering.hierarchical_clusterer import HierarchicalClusterer
            if isinstance(clustering_engine, HierarchicalClusterer):
                clustering_engine.create_dendrogram(embeddings)

        # Создание визуализации для DBSCAN
        if args.method == 'dbscan' and args.create_visualization:
            from clustering.dbscan_clusterer import DBSCANClusterer
            if isinstance(clustering_engine, DBSCANClusterer):
                clustering_engine.create_cluster_visualization(embeddings, labels)

        # Получение центров кластеров
        cluster_centers = clustering_engine.get_cluster_centers(embeddings, labels, clustering_model)

        # Вычисление схожестей
        similarities = clustering_engine.calculate_similarities(embeddings, cluster_centers, labels)

        # Извлечение ключевых слов
        topic_words = clustering_engine.extract_topic_keywords(
            valid_df[DB_CONFIG['text_column']], labels
        )

        # Подготовка результатов
        valid_df['cluster'] = labels
        valid_df['similarity'] = 0.0

        for cluster_id in np.unique(labels):
            if cluster_id == -1:  # Пропускаем шум
                continue
            cluster_mask = labels == cluster_id
            if cluster_id in similarities:  # Убедимся, что кластер есть в similarities
                valid_df.loc[cluster_mask, 'similarity'] = similarities[cluster_id]

        # Сохранение полных результатов в pkl, csv и json файлы
        results_manager.save_clustering_results(valid_df, embeddings, labels, topic_words, similarities,
                     cluster_centers, args.method, args.keyword_method)

        # Генерация отчетов с опциональными описаниями
        generate_descriptions = not args.do_descriptions
        # Если указана конкретная модель, обновляем конфиг
        if args.ollama_model:
            LLM_CONFIG["ollama"]["default_model"] = args.ollama_model

        report_generator.generate_all_reports(
            valid_df, topic_words, similarities, args.method, generate_descriptions,
            output_format=args.output_format
        )

        # Статистика
        stats = embedding_manager.get_collection_stats()
        method_info = clustering_engine.get_method_info()

        print(f"\n📊 Статистика ChromaDB: {stats}")
        print(f"🔧 Метод кластеризации: {method_info['name']}")
        print(f"📝 Описание: {method_info.get('description', '')}")

        if args.method == 'dbscan' and clustering_engine.optimal_eps:
            print(f"📏 Оптимальное eps: {clustering_engine.optimal_eps:.3f}")

        print(f"\n✅ Кластеризация завершена успешно!")
        print(f"📁 Результаты сохранены в: {RESULTS_DIR}")



    except Exception as e:
        print(f"❌ Ошибка: {e}")
        import traceback
        traceback.print_exc()
    finally:
        if 'embedding_manager' in locals():
            embedding_manager.cleanup()
        if 'clustering_engine' in locals() and hasattr(clustering_engine, 'cleanup'):
            clustering_engine.cleanup()

if __name__ == "__main__":
    main()