diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..989bad0 --- /dev/null +++ b/.env.example @@ -0,0 +1,10 @@ +# OpenCode API (DeepSeek V4 Flash Free) +OPENCODE_API_KEY=ваш-ключ-здесь +OPENCODE_URL=https://opencode.ai/zen/v1 + +# LM Studio (VLM для описаний PNG) +LMSTUDIO_URL=http://127.0.0.1:1234/v1 +LMSTUDIO_API_KEY=lm-studio + +# OpenAI (опционально) +# OPENAI_API_KEY=sk-... diff --git a/.gitignore b/.gitignore index 19ea8c9..6cbaf7e 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,17 @@ -output/*.png -output_123/*.png +# Generated outputs +output/ +output_123/ +lightrag_cache/ +vlm_descriptions.json +full_ocr_results.json + +# Python __pycache__/ *.pyc + +# Secrets +.env + +# OS .DS_Store Thumbs.db diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..703fd54 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,47 @@ +# Agent Guidelines for `opencode` + +## Project Overview +Python utilities for PDF processing and index building. Main scripts: +- `build_index.py` +- `process_pdf.py` +- `process_pdf_full.py` +- `process_any_pdf.py` + +## Repository Location +- **Local path:** `/Users/kirillblinov/development/opencode/OCR/opencode` +- **Remote (Gitea):** `https://gts.meratalk.online/keboss/opencode.git` + +## Git Rules + +### Branching +- Default branch: `master` +- Create feature/fix branches from `master` if needed. +- Prefer short-lived branches; merge back via PR or explicit request. + +### Commits +- Write concise, meaningful commit messages in English or Russian (match repo style). +- Stage only relevant files; never commit secrets or large binaries. +- Do **not** run `git push`, `git reset`, `git rebase`, or `git commit` unless explicitly asked. +- If asked to commit, inspect `git status` and `git diff` first, then confirm with the user before pushing. + +### Pull & Sync +- Before starting work, run `git pull origin master` to ensure the latest state. +- If there are local changes and the remote has moved, ask the user how to proceed (stash, merge, or reset). + +### Remotes +- Primary remote: `origin` → `https://gts.meratalk.online/keboss/opencode.git` +- Do not add new remotes without user confirmation. + +## Environment & Runtime +- This is a Python project; no `package.json` or `requirements.txt` is present. +- Use system Python or a virtual environment if dependencies are needed. +- Do not install packages globally without user confirmation. + +## File Handling +- **Ignore patterns (`.gitignore`):** `*.pdf`, `output/*.png`, `output_123/*.png`, `__pycache__/`, `*.pyc`, `.DS_Store`, `Thumbs.db` +- Do not commit generated outputs (PNGs, PDFs, caches). +- Keep source scripts clean and avoid hardcoding absolute paths when possible. + +## Communication +- If you need to modify repository configuration (`.git/config`, hooks, etc.), ask the user first. +- Always summarize what changed before and after any git operation. diff --git a/README_RAG.md b/README_RAG.md new file mode 100644 index 0000000..d3284ee --- /dev/null +++ b/README_RAG.md @@ -0,0 +1,208 @@ +# RAG / LightRAG для анализа чертежей + +Этот модуль добавляет **Retrieval-Augmented Generation (RAG)** поверх существующего OCR-pipeline. Он превращает распознанный текст из `full_ocr_results.json` в **граф знаний** и позволяет задавать вопросы по чертежам на естественном языке. + +## Архитектура + +``` +PDF -> PyMuPDF + RapidOCR -> full_ocr_results.json + | + v + rag_indexer.py (LightRAG) + | + v + lightrag_cache/ (граф знаний) + | + v + rag_query.py (вопросы) +``` + +**Почему LightRAG?** +- Строит **граф сущностей и отношений** (этажи, оси, размеры, помещения) +- Поддерживает 4 режима поиска: `naive`, `local`, `global`, `hybrid` +- Для чертежей рекомендуется **`hybrid`** — объединяет локальный поиск по соседним сущностям и глобальный контекст проекта + +## Установка + +```bash +# 1. Установите зависимости +pip install -r requirements.txt + +# 2. Выберите LLM backend: +``` + +### Вариант A: OpenAI (быстро, требует интернет) +```bash +export OPENAI_API_KEY="sk-..." +# или в Windows: set OPENAI_API_KEY=sk-... +``` + +### Вариант B: Ollama (локально, приватно, требует RAM/VRAM) +```bash +# Установите Ollama: https://ollama.com + +# Скачайте модели (пример) +ollama pull qwen2.5:14b # LLM для ответов +ollama pull nomic-embed-text # embeddings + +# Запустите сервер +ollama serve +``` + +### Вариант C: LM Studio (локальный GUI, подходит для VLM) +```bash +# Установите LM Studio: https://lmstudio.ai +# 1. Загрузите VLM модель (например, qwen3-vl-4b) и запустите Local Server +# 2. Убедитесь, что сервер доступен: http://127.0.0.1:1234/v1 + +# Переменные окружения (при необходимости) +export LMSTUDIO_URL="http://127.0.0.1:1234/v1" +export LMSTUDIO_API_KEY="lm-studio" +``` + +### Вариант D: OpenCode API (DeepSeek V4 Flash Free) +```bash +# Требуется API-ключ OpenCode (бесплатный тариф) +# 1. Получите ключ на https://opencode.ai +# 2. Экспортируйте переменные: + +export OPENCODE_API_KEY="ваш-ключ" +export OPENCODE_URL="https://opencode.ai/zen/v1" # опционально, уже задан по умолчанию +``` + +## Использование + +### Шаг 1: OCR (если еще не сделано) +```bash +python process_any_pdf.py "Кронштадтский 16-18 НК1_ОСК (v3).pdf" output_kronshtadt +``` +Результат: папка `output_kronshtadt/` с `full_ocr_results.json`. + +### Шаг 2: Построение индекса + +**OpenAI:** +```bash +python rag_indexer.py output_kronshtadt --backend openai --model gpt-4o-mini +``` + +**Ollama (локально):** +```bash +python rag_indexer.py output_kronshtadt --backend ollama --model qwen2.5:14b --embed-model nomic-embed-text +``` + +Результат: папка `output_kronshtadt/lightrag_cache/` с графом знаний. + +**LM Studio (VLM для описаний):** +```bash +# Сгенерировать VLM-описания PNG (локальная Qwen) +python vlm_describer.py output_kronshtadt --model qwen/qwen3-vl-4b + +# Построить индекс через OpenCode API (LLM) +python rag_indexer.py output_kronshtadt --backend opencode --vlm-desc +``` + +**OpenCode API (DeepSeek V4 Flash Free):** +```bash +# Построение индекса через бесплатный DeepSeek +python rag_indexer.py output_kronshtadt --backend opencode --model opencode/deepseek-v4-flash-free + +# С VLM-описаниями (локальная Qwen + бесплатный DeepSeek) +python rag_indexer.py output_kronshtadt --backend opencode --vlm-desc +``` + +### Шаг 3 (опционально): VLM-описания через LM Studio + +Если у вас в LM Studio загружена **VLM модель** (например, `qwen3-vl-4b`), можно сгенерировать текстовые описания PNG перед индексацией: + +```bash +python vlm_describer.py output_kronshtadt --model qwen/qwen3-vl-4b +``` + +Это создаст `output_kronshtadt/vlm_descriptions.json` — текстовое описание каждой страницы. При `--vlm-desc` в `rag_indexer.py` эти описания добавятся к OCR-тексту, давая RAG понимание пространственной компоновки. + +**Важно:** 4B VLM хороша для базовых описаний, но на сложных А0-чертежах может галлюцинировать. Для лучшего качества используйте 7B+ VLM (Qwen2-VL 7B, InternVL 8B). + +### Шаг 4: Запросы + +```bash +# Через OpenCode API (рекомендуется — бесплатно) +python rag_query.py output_kronshtadt "На каких страницах показан план 3-го этажа?" --backend opencode --mode hybrid + +python rag_query.py output_kronshtadt "Какие размеры указаны между осями А и Б?" --backend opencode --mode hybrid + +python rag_query.py output_kronshtadt "Какие квартиры есть на этаже 2?" --backend opencode --mode hybrid + +# Режим local — ищет соседей сущностей в графе +python rag_query.py output_kronshtadt "Что находится рядом с осью В?" --backend opencode --mode local + +# Режим global — общий контекст проекта +python rag_query.py output_kronshtadt "Опиши общую структуру здания по чертежам." --backend opencode --mode global +``` + +## Как это работает под капотом + +### 1. Формирование документов +`rag_indexer.py` читает `full_ocr_results.json` и для каждой страницы создает текстовый документ вида: + +``` +=== Чертеж: страница 5 из 120 === +Изображение: page_005.png + +--- Текстовый слой PDF --- +План 3-го этажа... + +--- Распознанный текст (OCR) --- + "3-й этаж" (confidence=0.95, bbox=[100, 200, 150, 220]) + "Ось А" (confidence=0.88, bbox=[300, 400, 350, 420]) + "5400 мм" (confidence=0.92, bbox=[500, 600, 580, 620]) + +--- Извлеченные сущности --- +Этажи: 3-й этаж +Оси: А +Размеры: 5400 мм +``` + +### 2. Построение графа (LightRAG) +LightRAG пропускает каждый документ через LLM и извлекает: +- **Entities**: `3-й этаж`, `Ось А`, `5400 мм` +- **Relationships**: `3-й этаж -> содержит -> Ось А`, `Ось А -> имеет размер -> 5400 мм` + +Это сохраняется в `lightrag_cache/` как граф + векторные эмбеддинги. + +### 3. Поиск ответа +При запросе `"Какие размеры между осями А и Б?"`: +- **Hybrid** = Local (ищет узлы "Ось А", "Ось Б" и их соседей) + Global (общий контекст "размеры между осями") +- LLM генерирует ответ на основе найденных фрагментов + +## Советы по использованию для чертежей + +1. **Качество OCR критично** — если RapidOCR пропускает цифры размеров, RAG не найдет их. Проверяйте `full_ocr_results.json`. +2. **Структурируйте сущности** — в `rag_indexer.py` уже встроен парсер этажей, осей, размеров. При необходимости расширьте регулярные выражения под ваши ГОСТ-овские обозначения. +3. **Страницы как узлы** — если чертеж очень большой (А0), можно разбить PNG на фрагменты и индексировать их отдельно. +4. **VLM для сложных схем** — если нужно понимать геометрию (а не только текст), добавьте описания от зрительной модели (Qwen2-VL, LLaVA) в текст документов перед индексацией. + +## Расширение: добавление VLM-описаний + +Если у вас есть GPU, можно улучшить понимание чертежей: + +```python +# Псевдокод — добавить в rag_indexer.py перед insert +from transformers import Qwen2VLForConditionalGeneration, AutoProcessor + +vlm = Qwen2VLForConditionalGeneration.from_pretrained("Qwen/Qwen2-VL-7B-Instruct") +processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct") + +desc = vlm_describe_image(page_image_path, "Опиши этот чертеж: этажи, оси, размеры, помещения.") +doc_text += f"\n--- Описание модели зрения ---\n{desc}\n" +await rag.ainsert(doc_text) +``` + +Это даст RAG понимание пространственной компоновки, даже если OCR плохо распознал текст. + +## Файлы + +| Файл | Назначение | +|------|-----------| +| `rag_indexer.py` | Построение индекса LightRAG из OCR-результатов | +| `rag_query.py` | Запросы к индексу | +| `requirements.txt` | Зависимости (lightrag-hku, openai, etc.) | diff --git a/process_any_pdf.py b/process_any_pdf.py index 8e376fd..7e6d9c3 100644 --- a/process_any_pdf.py +++ b/process_any_pdf.py @@ -10,7 +10,7 @@ import sys import json import fitz from pathlib import Path -from rapidocr import RapidOCR +from rapidocr_onnxruntime import RapidOCR # ------------------------------------------------------------------ # Параметры @@ -40,12 +40,13 @@ def process_pdf(pdf_path: Path, out_dir: Path): res = engine(img_path) ocr_lines = [] - if res and res.txts is not None: - for txt, box, score in zip(res.txts, res.boxes, res.scores): + if res and res[0] is not None: + for item in res[0]: + box, txt, score = item ocr_lines.append({ "text": txt, "confidence": float(score), - "bbox": box.tolist() if hasattr(box, "tolist") else box + "bbox": box }) all_pages.append({ @@ -74,8 +75,8 @@ def main(): pdf_file = sys.argv[1] out_name = sys.argv[2] if len(sys.argv) > 2 else f"output_{Path(pdf_file).stem}" - pdf_path = Path(r"D:\TEST docs") / pdf_file - out_dir = Path(r"D:\TEST docs") / out_name + pdf_path = Path(pdf_file) if Path(pdf_file).exists() else Path(".") / pdf_file + out_dir = Path(out_name) if not pdf_path.exists(): print(f"[ERR] Файл не найден: {pdf_path}") diff --git a/process_pdf.py b/process_pdf.py index 717d5a2..60c8c63 100644 --- a/process_pdf.py +++ b/process_pdf.py @@ -16,7 +16,7 @@ from pathlib import Path # 1. RapidOCR - лёгкий ONNX-OCR (входит в зависимости Docling) # ------------------------------------------------------------------ try: - from rapidocr import RapidOCR + from rapidocr_onnxruntime import RapidOCR engine = RapidOCR() print("[OK] RapidOCR загружен") except Exception as e: @@ -75,14 +75,17 @@ def run_ocr(image_paths: list[Path]): for img_path in image_paths: print(f" OCR: {img_path.name} ...") res = engine(img_path) - # res - RapidOCROutput s atributami txts, boxes, scores + # res - tuple: (results, timing) + # results: list of [box, text, confidence] entries = [] - for txt, box, score in zip(res.txts, res.boxes, res.scores): - entries.append({ - "text": txt, - "confidence": float(score), - "bbox": box.tolist() if hasattr(box, "tolist") else box - }) + if res and res[0] is not None: + for item in res[0]: + box, txt, score = item + entries.append({ + "text": txt, + "confidence": float(score), + "bbox": box + }) results[img_path.name] = entries print(f" naydeno {len(entries)} strok") return results diff --git a/process_pdf_full.py b/process_pdf_full.py index 97ce3ee..f4d55f3 100644 --- a/process_pdf_full.py +++ b/process_pdf_full.py @@ -10,7 +10,7 @@ import sys import json import fitz from pathlib import Path -from rapidocr import RapidOCR +from rapidocr_onnxruntime import RapidOCR # ------------------------------------------------------------------ # 1. Конфигурация @@ -44,12 +44,14 @@ def process_page(doc: fitz.Document, page_num: int, dpi: int = 300): # --- 2.3 OCR --- res = engine(img_path) ocr_lines = [] - for txt, box, score in zip(res.txts, res.boxes, res.scores): - ocr_lines.append({ - "text": txt, - "confidence": float(score), - "bbox": box.tolist() if hasattr(box, "tolist") else box - }) + if res and res[0] is not None: + for item in res[0]: + box, txt, score = item + ocr_lines.append({ + "text": txt, + "confidence": float(score), + "bbox": box + }) return { "page_number": page_num + 1, diff --git a/rag_indexer.py b/rag_indexer.py new file mode 100644 index 0000000..72f8c72 --- /dev/null +++ b/rag_indexer.py @@ -0,0 +1,352 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Построение LightRAG индекса из OCR-результатов чертежей. + +Использование: + # С OpenAI (требуется OPENAI_API_KEY) + python rag_indexer.py --backend openai + + # С локальной моделью через Ollama (требуется запущенный Ollama) + python rag_indexer.py --backend ollama --model qwen2.5:14b + + # Через LM Studio (локальный OpenAI-compatible сервер) + python rag_indexer.py --backend lmstudio --model qwen2.5:14b + + # С VLM-описаниями (предварительно запустить vlm_describer.py) + python rag_indexer.py --backend lmstudio --vlm-desc + +Результат: папка /lightrag_cache с графом знаний. +""" + +import os +import sys +import json +import asyncio +import argparse +from pathlib import Path +from typing import Optional + +from dotenv import load_dotenv +load_dotenv() + +# ------------------------------------------------------------------ +# LightRAG imports +# ------------------------------------------------------------------ +try: + from lightrag import LightRAG, QueryParam + from lightrag.utils import EmbeddingFunc +except ImportError as e: + print(f"[ERR] LightRAG не установлен: {e}") + print(" Установите: pip install lightrag-hku") + sys.exit(1) + +# ------------------------------------------------------------------ +# LLM backends +# ------------------------------------------------------------------ +def get_openai_backend(model: str = "gpt-4o-mini"): + """OpenAI backend (требует OPENAI_API_KEY).""" + from lightrag.llm import openai_complete_if_cache + from lightrag.llm import openai_embedding + + async def llm_func(prompt, system_prompt=None, history_messages=[], **kwargs): + return await openai_complete_if_cache( + model, prompt, + system_prompt=system_prompt, + history_messages=history_messages, + **kwargs + ) + + async def embed_func(texts: list[str]) -> list[list[float]]: + return await openai_embedding(texts, model="text-embedding-3-small") + + embed_cfg = EmbeddingFunc( + embedding_dim=1536, + max_token_size=8192, + func=embed_func, + ) + return llm_func, embed_cfg + + +def get_ollama_backend(model: str = "qwen2.5:14b", embed_model: str = "nomic-embed-text"): + """Ollama backend (требует запущенный ollama serve).""" + from lightrag.llm import ollama_model_complete, ollama_embedding + + async def llm_func(prompt, system_prompt=None, history_messages=[], **kwargs): + return await ollama_model_complete( + model, prompt, + system_prompt=system_prompt, + history_messages=history_messages, + **kwargs + ) + + async def embed_func(texts: list[str]) -> list[list[float]]: + return await ollama_embedding(texts, model=embed_model) + + # nomic-embed-text -> 768 dim, check your model + embed_cfg = EmbeddingFunc( + embedding_dim=768, + max_token_size=8192, + func=embed_func, + ) + return llm_func, embed_cfg + + +def get_lmstudio_backend(model: str = "qwen2.5:14b"): + """LM Studio backend (OpenAI-compatible API на http://127.0.0.1:1234/v1).""" + from lightrag.llm import openai_complete_if_cache + from sentence_transformers import SentenceTransformer + + base_url = os.environ.get("LMSTUDIO_URL", "http://127.0.0.1:1234/v1") + api_key = os.environ.get("LMSTUDIO_API_KEY", "lm-studio") + + async def llm_func(prompt, system_prompt=None, history_messages=[], **kwargs): + return await openai_complete_if_cache( + model, prompt, + system_prompt=system_prompt, + history_messages=history_messages, + api_key=api_key, + base_url=base_url, + **kwargs + ) + + # Для эмбеддингов используем локальную sentence-transformers + # (qwen3-vl-4b не умеет embeddings, поэтому нужна отдельная модель) + embed_model = SentenceTransformer("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2") + + async def embed_func(texts: list[str]) -> list[list[float]]: + import numpy as np + embeddings = embed_model.encode(texts, convert_to_numpy=True) + return embeddings.tolist() + + embed_cfg = EmbeddingFunc( + embedding_dim=384, + max_token_size=512, + func=embed_func, + ) + return llm_func, embed_cfg + + +def get_opencode_backend(model: str = "opencode/deepseek-v4-flash-free"): + """OpenCode backend (DeepSeek V4 Flash Free, OpenAI-compatible API).""" + from openai import AsyncOpenAI + from sentence_transformers import SentenceTransformer + + base_url = os.environ.get("OPENCODE_URL", "https://opencode.ai/zen/v1") + api_key = os.environ.get("OPENCODE_API_KEY", "") + client = AsyncOpenAI(base_url=base_url, api_key=api_key) + + async def llm_func(prompt, system_prompt=None, history_messages=[], **kwargs): + messages = [] + if system_prompt: + messages.append({"role": "system", "content": system_prompt}) + if history_messages: + messages.extend(history_messages) + messages.append({"role": "user", "content": prompt}) + + response = await client.chat.completions.create( + model=model, + messages=messages, + temperature=kwargs.get("temperature", 0.3), + max_tokens=kwargs.get("max_tokens", 1024), + ) + return response.choices[0].message.content + + # Эмбеддинги — локальные sentence-transformers + embed_model = SentenceTransformer("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2") + + async def embed_func(texts: list[str]) -> list[list[float]]: + import numpy as np + embeddings = embed_model.encode(texts, convert_to_numpy=True) + return embeddings + + embed_cfg = EmbeddingFunc( + embedding_dim=384, + max_token_size=512, + func=embed_func, + ) + return llm_func, embed_cfg + + +# ------------------------------------------------------------------ +# Подготовка текстов из OCR +# ------------------------------------------------------------------ +def format_page_document(page: dict, page_idx: int, total: int, vlm_descs: dict = None) -> str: + """Превращает данные одной страницы в текстовый документ для RAG.""" + lines = [] + lines.append(f"=== Чертеж: страница {page['page_number']} из {total} ===") + lines.append(f"Изображение: {page['image']}") + lines.append("") + + # VLM-описание изображения (если есть) + if vlm_descs and page["image"] in vlm_descs: + lines.append("--- Описание модели зрения (VLM) ---") + lines.append(vlm_descs[page["image"]]) + lines.append("") + + # Текстовый слой PDF (если есть) + pdf_text = page.get("pdf_text_layer", "").strip() + if pdf_text: + lines.append("--- Текстовый слой PDF ---") + lines.append(pdf_text[:2000]) # обрежем, чтобы не перегружать + lines.append("") + + # OCR результаты + ocr_lines = page.get("ocr_lines", []) + if ocr_lines: + lines.append("--- Распознанный текст (OCR) ---") + for entry in ocr_lines: + txt = entry["text"].strip() + conf = entry.get("confidence", 0.0) + bbox = entry.get("bbox", []) + lines.append(f' "{txt}" (confidence={conf:.2f}, bbox={bbox})') + lines.append("") + + # Извлечь и явно перечислить сущности (помогает LLM построить граф) + entities = extract_entities_from_page(ocr_lines) + if any(entities.values()): + lines.append("--- Извлеченные сущности ---") + if entities["floors"]: + lines.append(f"Этажи: {', '.join(entities['floors'])}") + if entities["axes"]: + lines.append(f"Оси: {', '.join(entities['axes'])}") + if entities["dimensions"]: + lines.append(f"Размеры: {', '.join(entities['dimensions'])}") + if entities["rooms"]: + lines.append(f"Помещения: {', '.join(entities['rooms'])}") + lines.append("") + + return "\n".join(lines) + + +def extract_entities_from_page(ocr_lines: list[dict]) -> dict: + """Извлекает сущности из OCR-строк страницы.""" + import re + floors, axes, dims, rooms = set(), set(), set(), set() + + floor_re = re.compile(r"(\d+)[-\s]?й\s*этаж", re.I) + axis_re = re.compile(r"\b([А-Я])\b") + num_axis_re = re.compile(r"\b(\d{1,2})\b") + dim_re = re.compile(r"\b(\d{3,5})\s*м?[мм]?\b") + room_re = re.compile(r"([Кк]вартира|[Кк]в\.?\s*\d+|[Пп]омещение\s*\d+)", re.I) + + for entry in ocr_lines: + t = entry["text"] + for m in floor_re.finditer(t): + floors.add(m.group(0)) + for m in axis_re.finditer(t): + axes.add(m.group(1)) + for m in num_axis_re.finditer(t): + v = m.group(1) + if v.isdigit() and 1 <= int(v) <= 50: + axes.add(v) + for m in dim_re.finditer(t): + v = m.group(1) + if 100 <= int(v) <= 20000: + dims.add(v + " мм") + for m in room_re.finditer(t): + rooms.add(m.group(0)) + + return { + "floors": sorted(floors), + "axes": sorted(axes), + "dimensions": sorted(dims), + "rooms": sorted(rooms), + } + + +def build_documents(ocr_path: Path, vlm_desc_path: Optional[Path] = None) -> list[str]: + """Собирает список текстовых документов из full_ocr_results.json и VLM-описаний.""" + data = json.loads(ocr_path.read_text(encoding="utf-8")) + pages = data["pages"] + + vlm_descs = {} + if vlm_desc_path and vlm_desc_path.exists(): + vlm_descs = json.loads(vlm_desc_path.read_text(encoding="utf-8")) + print(f" VLM-описания загружены: {len(vlm_descs)} шт.") + + docs = [] + total = len(pages) + + for i, page in enumerate(pages): + doc_text = format_page_document(page, i, total, vlm_descs) + docs.append(doc_text) + + return docs + + +# ------------------------------------------------------------------ +# Основной pipeline +# ------------------------------------------------------------------ +async def index_folder(folder: Path, backend: str, model: str, embed_model: Optional[str], use_vlm: bool = False): + ocr_path = folder / "full_ocr_results.json" + if not ocr_path.exists(): + print(f"[ERR] Не найден {ocr_path}") + sys.exit(1) + + cache_dir = folder / "lightrag_cache" + cache_dir.mkdir(exist_ok=True) + + vlm_desc_path = folder / "vlm_descriptions.json" if use_vlm else None + + print(f"[1/4] Загрузка OCR-данных из {ocr_path}") + docs = build_documents(ocr_path, vlm_desc_path) + print(f" Страниц: {len(docs)}") + + print(f"[2/4] Инициализация LightRAG (backend={backend}, model={model})") + if backend == "openai": + llm_func, embed_cfg = get_openai_backend(model) + elif backend == "ollama": + embed = embed_model or "nomic-embed-text" + llm_func, embed_cfg = get_ollama_backend(model, embed) + elif backend == "lmstudio": + llm_func, embed_cfg = get_lmstudio_backend(model) + elif backend == "opencode": + llm_func, embed_cfg = get_opencode_backend(model) + else: + print(f"[ERR] Неизвестный backend: {backend}") + sys.exit(1) + + rag = LightRAG( + working_dir=str(cache_dir), + llm_model_func=llm_func, + embedding_func=embed_cfg, + ) + + print(f"[2.5/4] Инициализация хранилищ...") + await rag.initialize_storages() + + print(f"[3/4] Вставка документов в индекс...") + for i, doc_text in enumerate(docs, 1): + print(f" [{i}/{len(docs)}] Страница {i}") + await rag.ainsert(doc_text) + + print(f"[4/4] Готово. Индекс сохранен в {cache_dir}") + print(f" Для запросов используйте: python rag_query.py \"{folder}\"") + + +def main(): + parser = argparse.ArgumentParser(description="Построение LightRAG индекса из OCR") + parser.add_argument("folder", help="Папка с full_ocr_results.json") + parser.add_argument("--backend", choices=["openai", "ollama", "lmstudio", "opencode"], default="openai", + help="LLM backend (default: openai)") + parser.add_argument("--model", default="gpt-4o-mini", + help="Название модели LLM (default: gpt-4o-mini, opencode: deepseek-v4-flash-free)") + parser.add_argument("--embed-model", default=None, + help="Модель эмбеддингов (только для Ollama, default: nomic-embed-text)") + parser.add_argument("--vlm-desc", action="store_true", + help="Использовать VLM-описания из vlm_descriptions.json") + args = parser.parse_args() + + folder = Path(args.folder) + model = args.model + if args.backend in ("ollama", "lmstudio") and model == "gpt-4o-mini": + model = "qwen2.5:14b" + if args.backend == "opencode" and model == "gpt-4o-mini": + model = "mimo-v2.5-free" + + asyncio.run(index_folder(folder, args.backend, model, args.embed_model, args.vlm_desc)) + + +if __name__ == "__main__": + main() diff --git a/rag_query.py b/rag_query.py new file mode 100644 index 0000000..867bc8a --- /dev/null +++ b/rag_query.py @@ -0,0 +1,193 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Запросы к LightRAG индексу по чертежам. + +Использование: + python rag_query.py "Какие размеры между осями А и Б?" + +Режимы поиска: + naive — прямой поиск по чанкам (без графа) + local — поиск по локальному графу соседей сущностей + global — поиск по глобальному контексту проекта + hybrid — комбинация local + global (рекомендуется для чертежей) +""" + +import os +import sys +import asyncio +import argparse +from pathlib import Path + +from dotenv import load_dotenv +load_dotenv() + +try: + from lightrag import LightRAG, QueryParam + from lightrag.utils import EmbeddingFunc +except ImportError as e: + print(f"[ERR] LightRAG не установлен: {e}") + sys.exit(1) + + +def get_openai_backend(model: str = "gpt-4o-mini"): + from lightrag.llm import openai_complete_if_cache, openai_embedding + + async def llm_func(prompt, system_prompt=None, history_messages=[], **kwargs): + return await openai_complete_if_cache( + model, prompt, system_prompt=system_prompt, + history_messages=history_messages, **kwargs + ) + + async def embed_func(texts: list[str]) -> list[list[float]]: + return await openai_embedding(texts, model="text-embedding-3-small") + + return llm_func, EmbeddingFunc(embedding_dim=1536, max_token_size=8192, func=embed_func) + + +def get_ollama_backend(model: str = "qwen2.5:14b", embed_model: str = "nomic-embed-text"): + from lightrag.llm import ollama_model_complete, ollama_embedding + + async def llm_func(prompt, system_prompt=None, history_messages=[], **kwargs): + return await ollama_model_complete( + model, prompt, system_prompt=system_prompt, + history_messages=history_messages, **kwargs + ) + + async def embed_func(texts: list[str]) -> list[list[float]]: + return await ollama_embedding(texts, model=embed_model) + + return llm_func, EmbeddingFunc(embedding_dim=768, max_token_size=8192, func=embed_func) + + +def get_lmstudio_backend(model: str = "qwen2.5:14b"): + """LM Studio backend.""" + from openai import AsyncOpenAI + from sentence_transformers import SentenceTransformer + + base_url = os.environ.get("LMSTUDIO_URL", "http://127.0.0.1:1234/v1") + api_key = os.environ.get("LMSTUDIO_API_KEY", "lm-studio") + client = AsyncOpenAI(base_url=base_url, api_key=api_key) + + async def llm_func(prompt, system_prompt=None, history_messages=[], **kwargs): + messages = [] + if system_prompt: + messages.append({"role": "system", "content": system_prompt}) + if history_messages: + messages.extend(history_messages) + messages.append({"role": "user", "content": prompt}) + + response = await client.chat.completions.create( + model=model, + messages=messages, + temperature=kwargs.get("temperature", 0.3), + max_tokens=kwargs.get("max_tokens", 1024), + ) + return response.choices[0].message.content + + embed_model = SentenceTransformer("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2") + + async def embed_func(texts: list[str]) -> list[list[float]]: + import numpy as np + return embed_model.encode(texts, convert_to_numpy=True) + + return llm_func, EmbeddingFunc(embedding_dim=384, max_token_size=512, func=embed_func) + + +def get_opencode_backend(model: str = "nemotron-3-super-free"): + """OpenCode backend (DeepSeek V4 Flash Free).""" + from openai import AsyncOpenAI + from sentence_transformers import SentenceTransformer + + base_url = os.environ.get("OPENCODE_URL", "https://opencode.ai/zen/v1") + api_key = os.environ.get("OPENCODE_API_KEY", "") + client = AsyncOpenAI(base_url=base_url, api_key=api_key) + + async def llm_func(prompt, system_prompt=None, history_messages=[], **kwargs): + messages = [] + if system_prompt: + messages.append({"role": "system", "content": system_prompt}) + if history_messages: + messages.extend(history_messages) + messages.append({"role": "user", "content": prompt}) + + response = await client.chat.completions.create( + model=model, + messages=messages, + temperature=kwargs.get("temperature", 0.3), + max_tokens=kwargs.get("max_tokens", 1024), + ) + return response.choices[0].message.content + + embed_model = SentenceTransformer("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2") + + async def embed_func(texts: list[str]) -> list[list[float]]: + import numpy as np + return embed_model.encode(texts, convert_to_numpy=True) + + return llm_func, EmbeddingFunc(embedding_dim=384, max_token_size=512, func=embed_func) + + +async def query_index(folder: Path, question: str, mode: str, backend: str, model: str, embed_model: str): + cache_dir = folder / "lightrag_cache" + if not cache_dir.exists(): + print(f"[ERR] Не найден индекс {cache_dir}. Сначала запустите rag_indexer.py") + sys.exit(1) + + print(f"[INIT] Загрузка индекса из {cache_dir} (backend={backend}, model={model})") + if backend == "openai": + llm_func, embed_cfg = get_openai_backend(model) + elif backend == "ollama": + llm_func, embed_cfg = get_ollama_backend(model, embed_model) + elif backend == "lmstudio": + llm_func, embed_cfg = get_lmstudio_backend(model) + elif backend == "opencode": + llm_func, embed_cfg = get_opencode_backend(model) + else: + print(f"[ERR] Неизвестный backend: {backend}") + sys.exit(1) + + rag = LightRAG( + working_dir=str(cache_dir), + llm_model_func=llm_func, + embedding_func=embed_cfg, + ) + + print(f"[INIT] Инициализация хранилищ...") + await rag.initialize_storages() + + print(f"[QUERY] Режим: {mode}") + print(f"[QUERY] Вопрос: {question}\n") + print("=" * 60) + + result = await rag.aquery(question, param=QueryParam(mode=mode)) + + print("=" * 60) + print("\n[RESULT] Ответ:") + print(result) + print("") + + +def main(): + parser = argparse.ArgumentParser(description="Запросы к LightRAG по чертежам") + parser.add_argument("folder", help="Папка с lightrag_cache") + parser.add_argument("question", help="Вопрос на естественном языке") + parser.add_argument("--mode", choices=["naive", "local", "global", "hybrid"], + default="hybrid", help="Режим поиска (default: hybrid)") + parser.add_argument("--backend", choices=["openai", "ollama", "lmstudio", "opencode"], default="openai") + parser.add_argument("--model", default="gpt-4o-mini") + parser.add_argument("--embed-model", default="nomic-embed-text") + args = parser.parse_args() + + folder = Path(args.folder) + model = args.model + if args.backend in ("ollama", "lmstudio") and model == "gpt-4o-mini": + model = "qwen2.5:14b" + if args.backend == "opencode" and model == "gpt-4o-mini": + model = "mimo-v2.5-free" + + asyncio.run(query_index(folder, args.question, args.mode, args.backend, model, args.embed_model)) + + +if __name__ == "__main__": + main() diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..9ff7898 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,21 @@ +# Core project dependencies +PyMuPDF>=1.24.0 +rapidocr-onnxruntime>=1.4.0 +docling>=2.0.0 + +# RAG stack +lightrag-hku>=1.0.0 +numpy>=1.24.0 + +# Optional: for local LLM via Ollama (install Ollama separately from https://ollama.com) +# ollama-python is pulled by lightrag-hku when needed + +# Optional: for OpenAI API (set OPENAI_API_KEY env var) +openai>=1.0.0 + +# Local embeddings (used with LM Studio backend) +sentence-transformers>=3.0.0 + +# Utilities +tqdm>=4.66.0 +python-dotenv>=1.0.0 diff --git a/run_pipeline.bat b/run_pipeline.bat new file mode 100644 index 0000000..f32ea0b --- /dev/null +++ b/run_pipeline.bat @@ -0,0 +1,53 @@ +@echo off +REM run_pipeline.bat — полный pipeline: OCR → VLM → RAG индекс +REM Требования: Python 3.11+, LM Studio запущен + +set PDF=%1 +if "%PDF%"=="" set PDF=123.pdf + +set OUT=%2 +if "%OUT%"=="" set OUT=output_%PDF:.pdf=% + +echo ============================================== +echo Pipeline: %PDF% → %OUT% +echo ============================================== + +REM Найти python3.11 +where python3.11 >nul 2>nul +if %errorlevel%==0 ( + set PYTHON=python3.11 +) else ( + where python >nul 2>nul + if %errorlevel%==0 ( + set PYTHON=python + ) else ( + echo [ERR] Python не найден + exit /b 1 + ) +) + +REM 1. OCR +echo. +echo [1/4] OCR: PDF → PNG + JSON +%PYTHON% process_any_pdf.py "%PDF%" "%OUT%" + +REM 2. VLM descriptions +echo. +echo [2/4] VLM: PNG → описания (LM Studio) +%PYTHON% vlm_describer.py "%OUT%" --model qwen/qwen3-vl-4b + +REM 3. RAG Index +echo. +echo [3/4] RAG индекс: JSON + VLM → граф знаний +if exist "%OUT%\lightrag_cache" rmdir /s /q "%OUT%\lightrag_cache" +%PYTHON% rag_indexer.py "%OUT%" --backend opencode --model nemotron-3-super-free --vlm-desc + +REM 4. Test query +echo. +echo [4/4] Тестовый запрос +%PYTHON% rag_query.py "%OUT%" "Какие оси есть в чертеже?" --backend opencode --mode hybrid + +echo. +echo ============================================== +echo Готово! Индекс в %OUT%\lightrag_cache +echo ============================================== diff --git a/run_pipeline.sh b/run_pipeline.sh new file mode 100755 index 0000000..b4b64f7 --- /dev/null +++ b/run_pipeline.sh @@ -0,0 +1,38 @@ +#!/bin/bash +# run_pipeline.sh — полный pipeline: OCR → VLM → RAG индекс + +set -e + +PDF="${1:-123.pdf}" +OUT="${2:-output_123}" +PYTHON="/opt/homebrew/bin/python3.11" + +echo "==============================================" +echo " Pipeline: $PDF → $OUT" +echo "==============================================" + +# 1. OCR +echo "" +echo "[1/4] OCR: PDF → PNG + JSON" +$PYTHON process_any_pdf.py "$PDF" "$OUT" + +# 2. VLM descriptions +echo "" +echo "[2/4] VLM: PNG → описания (LM Studio)" +$PYTHON vlm_describer.py "$OUT" --model qwen/qwen3-vl-4b + +# 3. RAG Index +echo "" +echo "[3/4] RAG индекс: JSON + VLM → граф знаний" +rm -rf "$OUT/lightrag_cache" +$PYTHON rag_indexer.py "$OUT" --backend opencode --model nemotron-3-super-free --vlm-desc + +# 4. Test query +echo "" +echo "[4/4] Тестовый запрос" +$PYTHON rag_query.py "$OUT" "Какие оси есть в чертеже?" --backend opencode --mode hybrid + +echo "" +echo "==============================================" +echo " Готово! Индекс в $OUT/lightrag_cache" +echo "==============================================" diff --git a/test_model.py b/test_model.py new file mode 100644 index 0000000..133c225 --- /dev/null +++ b/test_model.py @@ -0,0 +1,95 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Быстрая проверка модели OpenCode на способность извлекать сущности (как требует LightRAG). + +Использование: + python test_model.py + +Примеры: + python test_model.py deepseek-v4-flash-free + python test_model.py nemotron-3-super-free + python test_model.py minimax-m2.5-free +""" + +import os +import sys +import asyncio +from dotenv import load_dotenv +from openai import AsyncOpenAI + +load_dotenv() + +TEST_PROMPT = """ +Извлеки сущности и связи из текста. Ответь строго в формате: + +### Сущности: +1. entity_name (тип) + +### Связи: +1. entity1 -> relationship -> entity2 + +Текст: +Страница 5 чертежа. 3-й этаж здания. Оси А и Б. Размер 5400 мм между осями. Квартира 101 на этаже 3. +""" + +async def test_model(model: str): + base_url = os.environ.get("OPENCODE_URL", "https://opencode.ai/zen/v1") + api_key = os.environ.get("OPENCODE_API_KEY", "") + + if not api_key: + print("[ERR] OPENCODE_API_KEY не найден в .env") + sys.exit(1) + + print(f"[TEST] Модель: {model}") + print(f"[TEST] URL: {base_url}") + print(f"[TEST] Отправка запроса...\n") + + client = AsyncOpenAI(base_url=base_url, api_key=api_key) + + try: + response = await client.chat.completions.create( + model=model, + messages=[ + {"role": "system", "content": "Ты — анализатор чертежей. Извлекай структурированные сущности и связи."}, + {"role": "user", "content": TEST_PROMPT} + ], + temperature=0.3, + max_tokens=1024, + ) + + result = response.choices[0].message.content + print("=" * 60) + print("ОТВЕТ МОДЕЛИ:") + print("=" * 60) + print(result) + print("=" * 60) + + # Простая проверка + has_entities = "### Сущности" in result or "entity" in result.lower() or "сущности" in result.lower() + has_relations = "### Связи" in result or "relationship" in result.lower() or "связи" in result.lower() + + print(f"\n[РЕЗУЛЬТАТ]") + print(f" Найдены сущности: {'ДА' if has_entities else 'НЕТ'}") + print(f" Найдены связи: {'ДА' if has_relations else 'НЕТ'}") + + if has_entities and has_relations: + print(f"\n✅ Модель {model} ПОДХОДИТ для LightRAG") + else: + print(f"\n❌ Модель {model} НЕ подходит для LightRAG (не следует формату)") + print(" Попробуйте более мощную модель (GPT-4o, Claude, Llama 3.1 70B)") + + except Exception as e: + print(f"\n[ERR] Ошибка API: {e}") + print(" Возможно, модель недоступна или требует другого имени") + +if __name__ == "__main__": + if len(sys.argv) < 2: + print("Usage: python test_model.py ") + print("Examples:") + print(" python test_model.py deepseek-v4-flash-free") + print(" python test_model.py nemotron-3-super-free") + print(" python test_model.py minimax-m2.5-free") + sys.exit(1) + + asyncio.run(test_model(sys.argv[1])) diff --git a/vlm_describer.py b/vlm_describer.py new file mode 100644 index 0000000..174f446 --- /dev/null +++ b/vlm_describer.py @@ -0,0 +1,114 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Генерация текстовых описаний PNG-страниц через VLM в LM Studio. + +Требования: + - Запущен LM Studio с загруженной моделью (например, qwen3-vl-4b) + - Сервер: http://127.0.0.1:1234/v1 + +Использование: + python vlm_describer.py [--prompt "..."] [--model MODEL] + +Результат: /vlm_descriptions.json +""" + +import os +import sys +import json +import base64 +import argparse +from pathlib import Path +from openai import OpenAI + +# ------------------------------------------------------------------ +# Конфигурация LM Studio +# ------------------------------------------------------------------ +LMSTUDIO_URL = os.environ.get("LMSTUDIO_URL", "http://127.0.0.1:1234/v1") +LMSTUDIO_KEY = os.environ.get("LMSTUDIO_API_KEY", "lm-studio") + +DEFAULT_PROMPT = ( + "Опиши этот чертеж подробно. Укажи:\n" + "- Какой это этаж (если видно)\n" + "- Какие оси обозначены\n" + "- Какие размеры указаны\n" + "- Какие помещения/квартиры видны\n" + "- Общую компоновку и заметные детали.\n" + "Отвечай на русском языке." +) + +client = OpenAI(base_url=LMSTUDIO_URL, api_key=LMSTUDIO_KEY) + + +def encode_image(image_path: Path) -> str: + with open(image_path, "rb") as f: + return base64.b64encode(f.read()).decode("utf-8") + + +def describe_image(image_path: Path, model: str, prompt: str) -> str: + """Отправляет PNG в VLM и получает текстовое описание.""" + b64 = encode_image(image_path) + data_url = f"data:image/png;base64,{b64}" + + response = client.chat.completions.create( + model=model, + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": prompt}, + {"type": "image_url", "image_url": {"url": data_url}}, + ], + } + ], + temperature=0.3, + max_tokens=512, # 4B модель быстро устаёт, не гоним длину + ) + return response.choices[0].message.content.strip() + + +def process_folder(folder: Path, model: str, prompt: str): + """Обрабатывает все PNG в папке и сохраняет описания.""" + png_files = sorted(folder.glob("page_*.png")) + if not png_files: + print(f"[ERR] В папке {folder} не найдены page_*.png") + sys.exit(1) + + out_path = folder / "vlm_descriptions.json" + descriptions = {} + + print(f"[INFO] Найдено {len(png_files)} изображений") + print(f"[INFO] LM Studio: {LMSTUDIO_URL}") + print(f"[INFO] Модель: {model}\n") + + for i, png in enumerate(png_files, 1): + print(f"[{i}/{len(png_files)}] {png.name} ...", end=" ", flush=True) + try: + desc = describe_image(png, model, prompt) + descriptions[png.name] = desc + print(f"OK ({len(desc)} chars)") + except Exception as e: + print(f"ERR: {e}") + descriptions[png.name] = f"[ERROR] {e}" + + with open(out_path, "w", encoding="utf-8") as f: + json.dump(descriptions, f, ensure_ascii=False, indent=2) + + print(f"\n[OK] Сохранено: {out_path}") + + +def main(): + parser = argparse.ArgumentParser(description="VLM-описания PNG через LM Studio") + parser.add_argument("folder", help="Папка с page_*.png") + parser.add_argument("--model", default="qwen/qwen3-vl-4b", + help="Имя модели в LM Studio (default: qwen/qwen3-vl-4b)") + parser.add_argument("--prompt", default=DEFAULT_PROMPT, + help="Промпт для VLM") + args = parser.parse_args() + + folder = Path(args.folder) + process_folder(folder, args.model, args.prompt) + + +if __name__ == "__main__": + main()