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]))
|