c756a5766b
- Add rag_indexer.py: build LightRAG index from OCR with OpenCode API - Add rag_query.py: query the knowledge graph - Add vlm_describer.py: generate VLM descriptions via LM Studio - Add test_model.py: quick check for LightRAG-compatible models - Add run_pipeline.sh and run_pipeline.bat: full OCR → VLM → RAG pipeline - Fix rapidocr import (rapidocr_onnxruntime) - Fix process_any_pdf.py paths for cross-platform use - Add .env.example, README_RAG.md, AGENTS.md - Update .gitignore for outputs and secrets
96 lines
3.5 KiB
Python
96 lines
3.5 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
Быстрая проверка модели OpenCode на способность извлекать сущности (как требует LightRAG).
|
||
|
||
Использование:
|
||
python test_model.py <model_name>
|
||
|
||
Примеры:
|
||
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 <model_name>")
|
||
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]))
|