second_constantine
fcda2be4a9
Updated import paths to use direct module references instead of relative paths. Implemented dynamic benchmark discovery based on the contents of the tests/ directory, allowing for more flexible benchmark configuration without requiring code changes. This change improves maintainability and makes it easier to add new benchmarks.
270 lines
11 KiB
Python
Executable File
270 lines
11 KiB
Python
Executable File
#!/usr/bin/env python3
|
||
"""
|
||
Скрипт для генерации тестовых данных для бенчмарка AI с использованием Ollama.
|
||
|
||
Генерирует тесты через LLM для категорий:
|
||
- translation (переводы)
|
||
- summarization (пересказы)
|
||
- codegen (генерация кода)
|
||
|
||
Поддерживает валидацию generated тестов.
|
||
"""
|
||
|
||
import argparse
|
||
import json
|
||
import os
|
||
import sys
|
||
from pathlib import Path
|
||
from typing import Dict, List, Optional
|
||
|
||
# Добавляем путь к исходникам, чтобы импортировать ollama_client
|
||
sys.path.insert(0, str(Path(__file__).parent.parent))
|
||
|
||
from ollama_client import OllamaClient
|
||
from src.constants import TEST_SEPARATOR
|
||
|
||
def generate_translation_test(ollama: OllamaClient, model: str) -> Dict[str, str]:
|
||
"""Генерирует тест для перевода через LLM."""
|
||
# Генерируем английский текст
|
||
en_prompt = 'Generate a simple English sentence for translation test. The sentence should be clear and not too long (5-10 words). Example: "Hello, how are you today?"'
|
||
en_text = ollama.generate(model, en_prompt).strip()
|
||
|
||
# Генерируем перевод
|
||
ru_prompt = f"""Translate the following English sentence to Russian:
|
||
"{en_text}"
|
||
Provide only the translation, no additional text."""
|
||
ru_text = ollama.generate(model, ru_prompt).strip()
|
||
|
||
return {
|
||
"prompt": f"Translate the following English text to Russian: '{en_text}'",
|
||
"expected": ru_text
|
||
}
|
||
|
||
def generate_summarization_test(ollama: OllamaClient, model: str) -> Dict[str, str]:
|
||
"""Генерирует тест для пересказа через LLM."""
|
||
# Генерируем текст для пересказа
|
||
text_prompt = 'Generate a short text (3-5 sentences) for summarization test. The text should be about technology, science, or programming. Example: "Artificial intelligence is intelligence demonstrated by machines. It involves studying intelligent agents that perceive their environment and take actions to achieve goals."'
|
||
text = ollama.generate(model, text_prompt).strip()
|
||
|
||
# Генерируем пересказ
|
||
summary_prompt = f"""Summarize the following text in 1-2 sentences:
|
||
"{text}"
|
||
Provide only the summary, no additional text."""
|
||
summary = ollama.generate(model, summary_prompt).strip()
|
||
|
||
return {
|
||
"prompt": f"Summarize the following text in 1-2 sentences: '{text}'",
|
||
"expected": summary
|
||
}
|
||
|
||
def generate_codegen_test(ollama: OllamaClient, model: str) -> Dict[str, str]:
|
||
"""Генерирует тест для генерации кода через LLM."""
|
||
# Генерируем задание для кода
|
||
task_prompt = 'Generate a simple Python programming task. The task should be clear and specific, asking to write a function. Example: "Write a Python function that calculates the factorial of a number using recursion."'
|
||
task = ollama.generate(model, task_prompt).strip()
|
||
|
||
# Генерируем код
|
||
code_prompt = f"""Write Python code to solve the following task:
|
||
"{task}"
|
||
Provide only the code, no explanations or additional text."""
|
||
code = ollama.generate(model, code_prompt).strip()
|
||
|
||
return {
|
||
"prompt": task,
|
||
"expected": code
|
||
}
|
||
|
||
def generate_test(ollama: OllamaClient, model: str, category: str) -> Dict[str, str]:
|
||
"""Генерирует тест для указанной категории через LLM."""
|
||
if category == "translation":
|
||
return generate_translation_test(ollama, model)
|
||
elif category == "summarization":
|
||
return generate_summarization_test(ollama, model)
|
||
elif category == "codegen":
|
||
return generate_codegen_test(ollama, model)
|
||
else:
|
||
raise ValueError(f"Unknown category: {category}")
|
||
|
||
def validate_test(test_data: Dict[str, str]) -> bool:
|
||
"""Валидирует тестовые данные."""
|
||
if not isinstance(test_data, dict):
|
||
print("❌ Тест должен быть словарём (JSON объект)")
|
||
return False
|
||
|
||
if "prompt" not in test_data:
|
||
print("❌ Отсутствует поле 'prompt'")
|
||
return False
|
||
|
||
if "expected" not in test_data:
|
||
print("❌ Отсутствует поле 'expected'")
|
||
return False
|
||
|
||
if not isinstance(test_data["prompt"], str):
|
||
print("❌ Поле 'prompt' должно быть строкой")
|
||
return False
|
||
|
||
if not isinstance(test_data["expected"], str):
|
||
print("❌ Поле 'expected' должно быть строкой")
|
||
return False
|
||
|
||
if not test_data["prompt"].strip():
|
||
print("❌ Поле 'prompt' не может быть пустым")
|
||
return False
|
||
|
||
if not test_data["expected"].strip():
|
||
print("❌ Поле 'expected' не может быть пустым")
|
||
return False
|
||
|
||
return True
|
||
|
||
def validate_all_tests(test_dir: str) -> None:
|
||
"""Валидирует все тесты в указанной директории."""
|
||
test_dir_path = Path(test_dir)
|
||
if not test_dir_path.exists():
|
||
print(f"❌ Директория {test_dir} не существует")
|
||
return
|
||
|
||
valid_count = 0
|
||
invalid_count = 0
|
||
|
||
for txt_file in test_dir_path.glob("*.txt"):
|
||
try:
|
||
with open(txt_file, "r", encoding="utf-8") as f:
|
||
content = f.read()
|
||
|
||
# Разделяем по разделителю
|
||
parts = content.split(TEST_SEPARATOR, 1)
|
||
if len(parts) != 2:
|
||
invalid_count += 1
|
||
print(f"❌ {txt_file.name} - некорректный формат (отсутствует разделитель)")
|
||
continue
|
||
|
||
prompt, expected = parts
|
||
test_data = {
|
||
"prompt": prompt,
|
||
"expected": expected
|
||
}
|
||
|
||
if validate_test(test_data):
|
||
valid_count += 1
|
||
print(f"✅ {txt_file.name} - валиден")
|
||
else:
|
||
invalid_count += 1
|
||
print(f"❌ {txt_file.name} - не валиден")
|
||
except Exception as e:
|
||
invalid_count += 1
|
||
print(f"❌ {txt_file.name} - ошибка: {str(e)}")
|
||
|
||
print(f"\nРезультаты валидации:")
|
||
print(f"Валидных тестов: {valid_count}")
|
||
print(f"Невалидных тестов: {invalid_count}")
|
||
print(f"Всего тестов: {valid_count + invalid_count}")
|
||
|
||
def generate_tests(ollama: OllamaClient, model: str, count: int, category: str, output_dir: str) -> None:
|
||
"""Генерирует указанное количество тестов через LLM."""
|
||
if category not in ["translation", "summarization", "codegen", "all"]:
|
||
print(f"❌ Неизвестная категория: {category}")
|
||
return
|
||
|
||
categories = [category] if category != "all" else ["translation", "summarization", "codegen"]
|
||
|
||
for cat in categories:
|
||
cat_dir = Path(output_dir) / cat
|
||
cat_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
for i in range(1, count + 1):
|
||
# Проверяем, существует ли уже тест с таким номером
|
||
test_num = 1
|
||
while True:
|
||
test_file = cat_dir / f"test{test_num}.txt"
|
||
if not test_file.exists():
|
||
break
|
||
test_num += 1
|
||
|
||
print(f"🤖 Генерирую тест {cat}/test{test_num}.txt...")
|
||
|
||
# Генерируем тест через LLM
|
||
test_data = generate_test(ollama, model, cat)
|
||
|
||
# Валидируем generated тест
|
||
if not validate_test(test_data):
|
||
print(f"❌ Сгенерирован невалидный тест для {cat}, тест номер {test_num}")
|
||
continue
|
||
|
||
# Сохраняем тест в формате TXT с разделителем
|
||
with open(test_file, "w", encoding="utf-8") as f:
|
||
f.write(f"{test_data['prompt']}{TEST_SEPARATOR}{test_data['expected']}")
|
||
|
||
print(f"✅ Создан тест {cat}/test{test_num}.txt")
|
||
|
||
def main():
|
||
"""Основная функция скрипта."""
|
||
parser = argparse.ArgumentParser(
|
||
description="Генератор тестовых данных для AI бенчмарка с использованием Ollama",
|
||
epilog="Примеры использования:\n"
|
||
" python scripts/generate_tests.py --count 2 --category translation --model second_constantine/t-lite-it-1.0:7b --ollama-url http://10.0.0.4:11434\n"
|
||
" python scripts/generate_tests.py --category all --model second_constantine/t-lite-it-1.0:7b --ollama-url http://10.0.0.4:11434\n"
|
||
" python scripts/generate_tests.py --validate tests/translation"
|
||
)
|
||
parser.add_argument(
|
||
"--count",
|
||
type=int,
|
||
default=1,
|
||
help="Количество тестов для генерации (по умолчанию: 1)"
|
||
)
|
||
parser.add_argument(
|
||
"--category",
|
||
type=str,
|
||
default="all",
|
||
choices=["translation", "summarization", "codegen", "all"],
|
||
help="Категория тестов (translation, summarization, codegen, или all) (по умолчанию: all)"
|
||
)
|
||
parser.add_argument(
|
||
"--model",
|
||
type=str,
|
||
help="Название модели для генерации тестов (требуется при генерации тестов)"
|
||
)
|
||
parser.add_argument(
|
||
"--ollama-url",
|
||
type=str,
|
||
help="URL подключения к Ollama серверу (требуется при генерации тестов)"
|
||
)
|
||
parser.add_argument(
|
||
"--validate",
|
||
type=str,
|
||
help="Валидировать тесты в указанной директории (например: tests/translation)"
|
||
)
|
||
|
||
args = parser.parse_args()
|
||
|
||
if args.validate:
|
||
print(f"🔍 Начинаю валидацию тестов в {args.validate}")
|
||
validate_all_tests(args.validate)
|
||
print("\n✨ Готово!")
|
||
sys.exit(0)
|
||
|
||
# Проверяем обязательные параметры для генерации тестов
|
||
if not args.model:
|
||
print("❌ Ошибка: параметр --model обязателен при генерации тестов")
|
||
sys.exit(1)
|
||
|
||
if not args.ollama_url:
|
||
print("❌ Ошибка: параметр --ollama-url обязателен при генерации тестов")
|
||
sys.exit(1)
|
||
|
||
print(f"🤖 Подключаюсь к Ollama серверу: {args.ollama_url}")
|
||
print(f"📝 Генерирую {args.count} тест(ов) для категории: {args.category}")
|
||
print(f"🎯 Используемая модель: {args.model}")
|
||
|
||
try:
|
||
ollama = OllamaClient(args.ollama_url)
|
||
generate_tests(ollama, args.model, args.count, args.category, "tests")
|
||
except Exception as e:
|
||
print(f"❌ Ошибка при генерации тестов: {e}")
|
||
sys.exit(1)
|
||
|
||
print("\n✨ Готово!")
|
||
|
||
if __name__ == "__main__":
|
||
main()
|