tests/make_clust.py - common_agents_notes - IvanMysin/Topics

tests/make_clust.py
158 строк7 KB
ivan
first_commit
24 сен 2025, 18:38
24 сен 2025, 18:380ac5e09
import sqlite3
import pandas as pd
import numpy as np
from sentence_transformers import SentenceTransformer
from bertopic import BERTopic
from umap import UMAP
from hdbscan import HDBSCAN
from sklearn.feature_extraction.text import CountVectorizer
import pickle
from tqdm.auto import tqdm

# Конфигурация
DB_PATH = "../data/hippocampal_rhythms_articles.db"  # "documents.db"  # Путь к вашей SQLite базе
TABLE_NAME = "Articles"  # Название таблицы с документами
TEXT_COLUMN = "full_text"      # Название колонки с текстом
ID_COLUMN = "id"          # Название колонки с ID документа
OUTPUT_FILE = "../data/topic_modeling_results.pkl"  # Файл для сохранения результатов

def load_documents_from_sqlite(db_path, table_name, text_column, id_column):
    """Загружает документы из SQLite базы данных"""
    print("Загрузка документов из SQLite...")
    try:
        conn = sqlite3.connect(db_path)
        query = f"SELECT {id_column}, {text_column} FROM {table_name}"
        df = pd.read_sql_query(query, conn)
        conn.close()

        # Удаляем пустые документы
        df = df.dropna(subset=[text_column])
        df = df[df[text_column].str.strip().astype(bool)]

        print(f"Загружено {len(df)} документов")
        return df
    except Exception as e:
        print(f"Ошибка при загрузке данных: {e}")
        return None

def compute_embeddings(documents, model_name="allenai/specter2_base"):
    """Вычисляет эмбеддинги для документов с помощью SPECTER2"""
    print(f"Загрузка модели {model_name}...")
    model = SentenceTransformer(model_name)

    print("Вычисление эмбеддингов...")
    # Включаем прогресс-бар для больших коллекций
    embeddings = model.encode(
        documents.tolist(),
        show_progress_bar=True,
        batch_size=32,  # Можно настроить в зависимости от GPU/CPU
        convert_to_numpy=True
    )

    print(f"Эмбеддинги вычислены. Размерность: {embeddings.shape}")
    return embeddings

def create_topic_model():
    """Создает и настраивает модель BERTopic"""
    # Уменьшаем размерность эмбеддингов
    umap_model = UMAP(
        n_components=5,           # Низкая размерность для научных текстов
        n_neighbors=15,           # Количество соседей
        min_dist=0.1,             # Минимальное расстояние
        metric='cosine',          # Метрика расстояния
        random_state=42
    )

    # Кластеризация
    hdbscan_model = HDBSCAN(
        min_cluster_size=10,      # Минимальный размер кластера
        min_samples=5,            # Минимальное количество samples
        metric='euclidean',       # Метрика расстояния
        cluster_selection_epsilon=0.1,  # Параметр для контроля плотности
        prediction_data=True      # Включаем предсказание для новых данных
    )

    # Векторизатор для обработки текста
    vectorizer_model = CountVectorizer(
        stop_words="english",     # Удаляем стоп-слова
        ngram_range=(1, 2),       # Учитываем униграммы и биграммы
        min_df=2,                 # Минимальная частота термина
        max_df=0.95               # Максимальная частота термина
    )

    # Создаем BERTopic модель
    topic_model = BERTopic(
        umap_model=umap_model,
        hdbscan_model=hdbscan_model,
        vectorizer_model=vectorizer_model,
        language="english",       # Язык для стоп-слов
        calculate_probabilities=True,  # Включаем вероятностные предсказания
        verbose=True
    )

    return topic_model

def save_results(topic_model, df, embeddings, output_file):
    """Сохраняет все результаты в файл"""
    results = {
        'topic_model': topic_model,
        'documents_df': df,
        'embeddings': embeddings,
        'topic_info': topic_model.get_topic_info(),
        'document_topics': topic_model.get_document_info(df[TEXT_COLUMN].tolist())
    }

    with open(output_file, 'wb') as f:
        pickle.dump(results, f)

    print(f"Результаты сохранены в {output_file}")

def main():
    # 1. Загрузка документов
    df = load_documents_from_sqlite(DB_PATH, TABLE_NAME, TEXT_COLUMN, ID_COLUMN)
    if df is None or len(df) == 0:
        return

    # 2. Вычисление эмбеддингов
    embeddings = compute_embeddings(df[TEXT_COLUMN])

    # 3. Создание и обучение topic model
    print("Создание и обучение модели BERTopic...")
    topic_model = create_topic_model()

    # Преобразуем тексты в список для BERTopic
    documents_list = df[TEXT_COLUMN].tolist()

    # Обучаем модель (fit_transform)
    topics, probabilities = topic_model.fit_transform(documents_list, embeddings)

    # 4. Сохранение результатов
    save_results(topic_model, df, embeddings, OUTPUT_FILE)

    # 5. Вывод основной информации
    print("\n" + "="*50)
    print("РЕЗУЛЬТАТЫ КЛАСТЕРИЗАЦИИ")
    print("="*50)

    # Информация о темах
    topic_info = topic_model.get_topic_info()
    print(f"Обнаружено тем: {len(topic_info) - 1}")  # -1 потому что тема -1 это outliers
    print(f"Документов без четкой темы (outliers): {len(topic_info[topic_info['Topic'] == -1])}")

    # Показываем топ-5 тем
    print("\nТоп-5 тем:")
    for topic_num in topic_info['Topic'].head(6):  # 6 потому что включая -1
        if topic_num != -1:
            print(f"\nТема {topic_num}:")
            words = topic_model.get_topic(topic_num)
            for word, score in words[:5]:  # Топ-5 слов для темы
                print(f"  {word} (score: {score:.3f})")

    # Сохранение тем в CSV для удобства просмотра
    doc_info = topic_model.get_document_info(documents_list)
    doc_info[ID_COLUMN] = df[ID_COLUMN].values
    doc_info.to_csv("document_topics.csv", index=False)
    print(f"\nДетальная информация по документам сохранена в document_topics.csv")

if __name__ == "__main__":
    main()