KG-Library

Библиотека для автоматического построения графов знаний (Knowledge Graphs) из текстовых, аудио и видео данных с использованием методов машинного обучения и обработки естественного языка.

📋 Описание

KG-Library — это комплексная система для извлечения, фильтрации и построения графов знаний из различных источников информации. Проект предназначен для работы с литературными произведениями и позволяет автоматически создавать структурированные представления знаний о книгах, их авторах, персонажах, местах действия и взаимосвязях между ними.

Основные возможности

• Извлечение триплетов из текста с использованием модели mREBEL (Multilingual Relation Extraction)
• Обработка мультимедиа: поддержка текстовых файлов, аудио (MP3, WAV) и видео (MP4, AVI, MOV, MKV)
• Обогащение данных из Wikidata и OpenLibrary
• Фильтрация триплетов с помощью обученных Graph Neural Networks
• Предсказание связей между сущностями в графе
• Обучение и дообучение моделей на пользовательских данных
• Интеграция с Neo4j для хранения и визуализации графов знаний
• Визуализация метрик обучения через TensorBoard

🛠 Технологии

Основные библиотеки и фреймворки

• Python 3.11 — основной язык программирования
• PyTorch (^2.1.0) — фреймворк для глубокого обучения
• PyTorch Geometric (^2.6.1) — работа с графовыми нейронными сетями
• Transformers (^4.31.0) — модели NLP (mREBEL)
• LangChain (^0.2.11) — работа с языковыми моделями
• spaCy (^3.5.4) — обработка естественного языка
• Neo4j (5.28.1) — графовая база данных
• OpenAI Whisper (^20240930) — распознавание речи из аудио/видео
• Sentence Transformers (^2.2.0) — генерация эмбеддингов
• scikit-learn (1.6.1) — метрики и оценка моделей
• TensorBoard (^2.13.0) — визуализация процесса обучения

Дополнительные инструменты

• Poetry — управление зависимостями
• Docker & Docker Compose — контейнеризация
• CUDA 12.4.1 — поддержка GPU для ускорения вычислений
• BeautifulSoup4 — парсинг HTML (для OpenLibrary)
• pandas (2.2.3) — работа с данными
• seaborn (0.13.2) — визуализация данных

Архитектура

• CompGCN (Composition-based Graph Convolutional Network) — архитектура GNN для работы с гетерогенными графами
• mREBEL — модель для извлечения отношений из текста
• Wikidata SPARQL API — получение структурированных данных о книгах и авторах
• OpenLibrary API — дополнительная информация о персонажах

Установка

Требования

• Python 3.11
• CUDA 12.4+ (опционально, для GPU)
• Docker и Docker Compose (для контейнеризации)
• Neo4j (через Docker или отдельная установка)

Установка через Docker (рекомендуется)

1. Клонируйте репозиторий:

git clone <repository-url>
cd KG_Library

2. Создайте файл .env на основе примера:

NEO4J_AUTH=neo4j/your_password
NEO4J_URI=neo4j://neo4j:7687
EMAIL=your_email@example.com

3. Запустите через Docker Compose:

docker-compose up --build

Локальная установка

1. Установите Poetry (если еще не установлен):

curl -sSL https://install.python-poetry.org | python3 -

2. Установите зависимости:

poetry install

3. Активируйте виртуальное окружение:

poetry shell

4. Загрузите модели spaCy:

python -m spacy download en_core_web_lg
python -m coreferee install en

5. Настройте переменные окружения в файле .env

Использование

Обработка текстового файла

python -m kg_library.main --input data/input/book.txt --model models/model.pt --confidence 0.65

Обработка аудио файла

python -m kg_library.main --input data/input/audio.mp3 --model models/model.pt --confidence 0.65

Обработка видео файла

python -m kg_library.main --input data/input/video.mp4 --model models/model.pt --link-prediction

Обучение модели с нуля

python -m kg_library.main --learn --size-dataset 200 --finetune

Дообучение существующей модели

python -m kg_library.main --input data/input/new_book.txt --finetune --model models/base_model.pt

Параметры командной строки

• --input — путь к входному файлу (текст, аудио или видео)
• --model — путь к модели (.pt или .zip файл)
• --output — путь для сохранения обновленной модели (по умолчанию: models/kg_model_updated.zip)
• --confidence — порог уверенности для фильтрации триплетов (0-1, по умолчанию: 0.65)
• --link-prediction — включить предсказание связей внутри графа
• --learn — обучить модель с нуля
• --finetune — дообучить существующую модель
• --size-dataset — размер датасета для обучения (по умолчанию: 200)
• --graph-path — путь к существующему графу для обучения (.json)
• --load-triplets — загрузить сохраненные триплеты из JSON
• --neo4j — загрузить граф из Neo4j
• --no-neo4j — не сохранять граф в Neo4j
• --no-save — не сохранять обновленную модель

Архитектура проекта

kg_library/
├── main.py                 # Точка входа CLI
├── AppFacade.py            # Основной фасад приложения
├── config.py               # Конфигурация
├── db.py                   # Подключение к Neo4j
├── common/                 # Общие компоненты
│   ├── GraphData.py       # Структура графа знаний
│   ├── NodeData.py        # Узлы графа
│   ├── EdgeData.py        # Рёбра графа
│   ├── GraphJSON.py       # Сериализация графов
│   └── BaseInformationExtractor.py  # Извлечение данных из Wikidata
├── models/                 # Модели машинного обучения
│   ├── entity_extraction/
│   │   └── TripletExtractor.py  # Извлечение триплетов (mREBEL)
│   ├── training/
│   │   ├── GraphNN.py          # Graph Neural Network
│   │   ├── CompGCNConv.py     # CompGCN слой
│   │   ├── GraphTrainer.py    # Обучение модели
│   │   ├── EmbeddingPreprocessor.py  # Предобработка данных
│   │   └── DataLoader.py      # Загрузка данных
│   └── evaluation/
│       └── TripletEvaluator.py  # Оценка триплетов
└── utils/                  # Утилиты
    ├── AudioProcessor.py   # Обработка аудио (Whisper)
    ├── VideoProcessor.py   # Обработка видео
    ├── Preprocessing.py    # Предобработка текста
    └── LogVisualizer.py    # Визуализация логов

Процесс работы

1. Извлечение триплетов: Текст обрабатывается моделью mREBEL для извлечения триплетов (сущность-отношение-сущность)
2. Обогащение данных: Система запрашивает дополнительную информацию из Wikidata и OpenLibrary
3. Фильтрация: Обученная GNN модель оценивает качество триплетов и фильтрует их по порогу уверенности
4. Построение графа: Отфильтрованные триплеты объединяются в граф знаний
5. Предсказание связей: Опционально система может предсказывать новые связи между сущностями
6. Сохранение: Граф сохраняется в Neo4j и/или в JSON файл

Обучение модели

Модель обучается на датасете книг (kingkangkr/book_summary_dataset), извлекая триплеты из аннотаций и обогащая их данными из Wikidata. Процесс обучения включает:

• Предобработку графа и создание эмбеддингов
• Обучение CompGCN модели для предсказания связей
• Валидацию и раннюю остановку
• Логирование метрик в TensorBoard

Для просмотра метрик обучения:

tensorboard --logdir=runs

Работа с Neo4j

Проект интегрирован с Neo4j для хранения и визуализации графов знаний. После обработки данных граф автоматически сохраняется в базу данных (если не указан флаг --no-neo4j).

Доступ к Neo4j Browser: http://localhost:7474

Форматы данных

Поддерживаемые входные форматы

• Текст: .txt, .md, .rst, .tex
• Аудио: .mp3, .wav, .ogg, .flac
• Видео: .mp4, .avi, .mov, .mkv

Выходные форматы

• JSON — сериализованный граф знаний
• Neo4j — граф в базе данных
• PyTorch — обученная модель (.pt или .zip)

Конфигурация

Основные параметры конфигурации находятся в kg_library/config.py:

• EMBEDDING_DIM — размерность эмбеддингов (по умолчанию: 64)
• HIDDEN_DIM — размерность скрытых слоёв (по умолчанию: 64)
• BATCHSIZE — размер батча (по умолчанию: 128)

Тестирование

Тесты находятся в директории kg_library/tests/:

python -m pytest kg_library/tests/

Автор

bleeeana (bobkov.v03@icloud.com)

Благодарности

• Babelscape за модель mREBEL
• OpenAI за Whisper
• Сообщество PyTorch Geometric
• Wikidata и OpenLibrary за открытые данные

---

Версия: 0.1.0