feat: add MongoDB test generation and update dependencies
- Added pymongo==3.13.0 to requirements.txt for MongoDB connectivity - Implemented generate_summarization_from_mongo.py script to generate summarization tests from MongoDB - Updated run.sh to support 'gen-mongo' command for MongoDB test generation - Enhanced scripts/README.md with documentation for new MongoDB functionality - Improved help text in run.sh to clarify available commands and usage examples ``` This commit adds MongoDB integration for test generation and updates the documentation and scripts accordingly.
This commit is contained in:
@@ -1,86 +1,69 @@
|
||||
# Скрипты для генерации тестов
|
||||
|
||||
Эта директория содержит утилиты для автоматизированной генерации и валидации тестовых данных с использованием Ollama.
|
||||
Эта директория содержит скрипты для генерации тестовых данных для AI бенчмарка.
|
||||
|
||||
## generate_tests.py
|
||||
## Доступные скрипты
|
||||
|
||||
Скрипт для генерации тестовых данных для AI бенчмарка через LLM.
|
||||
### 1. `generate_tests.py`
|
||||
|
||||
### Использование
|
||||
Скрипт для генерации тестовых данных через LLM (Ollama).
|
||||
|
||||
**Функциональность:**
|
||||
- Генерация тестов для переводов (translation)
|
||||
- Генерация тестов для пересказов (summarization)
|
||||
- Генерация тестов для генерации кода (codegen)
|
||||
- Валидация generated тестов
|
||||
|
||||
**Использование:**
|
||||
```bash
|
||||
# Генерировать 2 теста для перевода через Ollama
|
||||
python scripts/generate_tests.py --count 2 --category translation --model llama3 --ollama-url http://localhost:11434
|
||||
|
||||
# Генерировать 1 тест для каждого типа
|
||||
python scripts/generate_tests.py --count 1 --category all --model llama3 --ollama-url http://localhost:11434
|
||||
|
||||
# Генерировать 3 теста для пересказов
|
||||
python scripts/generate_tests.py --count 3 --category summarization --model llama3 --ollama-url http://localhost:11434
|
||||
|
||||
# Валидировать существующие тесты
|
||||
python scripts/generate_tests.py --validate tests/translation
|
||||
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
|
||||
```
|
||||
|
||||
### Аргументы
|
||||
|
||||
**Параметры:**
|
||||
- `--count`: Количество тестов для генерации (по умолчанию: 1)
|
||||
- `--category`: Категория тестов (translation, summarization, codegen, или all) (по умолчанию: all)
|
||||
- `--model`: Название модели для генерации тестов (обязательный параметр, например: llama3)
|
||||
- `--ollama-url`: URL подключения к Ollama серверу (обязательный параметр, например: http://localhost:11434)
|
||||
- `--model`: Название модели для генерации тестов (обязательный параметр)
|
||||
- `--ollama-url`: URL подключения к Ollama серверу (обязательный параметр)
|
||||
- `--validate`: Валидировать тесты в указанной директории
|
||||
|
||||
### Поддерживаемые категории
|
||||
### 2. `generate_summarization_from_mongo.py`
|
||||
|
||||
1. **translation** - тесты переводов с английского на русский (LLM генерирует английский текст и его перевод)
|
||||
2. **summarization** - тесты пересказов текстов (LLM генерирует текст и его пересказ)
|
||||
3. **codegen** - тесты генерации Python кода (LLM генерирует задание и код)
|
||||
Скрипт для генерации тестов пересказов из MongoDB.
|
||||
|
||||
### Как работает генерация
|
||||
**Функциональность:**
|
||||
- Извлекает текст статьи из коллекции `rssNotification` (поле `.meta.topicContent`)
|
||||
- Генерирует тестовые данные в формате JSON для бенчмарка AI
|
||||
- Валидирует generated тесты
|
||||
|
||||
Скрипт использует LLM для динамической генерации тестов:
|
||||
- **Translation**: LLM создает английский текст, затем переводит его на русский
|
||||
- **Summarization**: LLM генерирует текст о технологиях, затем создает его пересказ
|
||||
- **Codegen**: LLM формулирует задачу по программированию, затем пишет решение
|
||||
**Использование:**
|
||||
```bash
|
||||
python scripts/generate_summarization_from_mongo.py --record-id 507f1f77bcf86cd799439011
|
||||
```
|
||||
|
||||
### Примеры generated тестов
|
||||
**Параметры:**
|
||||
- `--record-id`: ID записи в MongoDB (обязательный параметр)
|
||||
- `--output-dir`: Директория для сохранения generated тестов (по умолчанию: tests/summarization)
|
||||
|
||||
#### Translation
|
||||
**Требования:**
|
||||
- Доступ к MongoDB кластеру (10.0.0.3, 10.0.0.4, 10.0.0.5)
|
||||
- Установленный пакет `pymongo` (автоматически устанавливается при первом запуске)
|
||||
|
||||
**Формат generated тестов:**
|
||||
```json
|
||||
{
|
||||
"prompt": "Translate the following English text to Russian: 'Hello, how are you today?'",
|
||||
"expected": "Привет, как дела сегодня?"
|
||||
"prompt": "Summarize the following text in 1-2 sentences: 'Текст статьи из MongoDB'",
|
||||
"expected": ""
|
||||
}
|
||||
```
|
||||
|
||||
#### Summarization
|
||||
```json
|
||||
{
|
||||
"prompt": "Summarize the following text in 1-2 sentences: 'The quick brown fox jumps over the lazy dog...'",
|
||||
"expected": "A quick fox jumps over a lazy dog, surprising it. The fox keeps running while the dog barks."
|
||||
}
|
||||
**Примечание:** Поле "expected" будет пустым, так как ожидаемый результат нужно будет сгенерировать отдельно через LLM или вручную.
|
||||
|
||||
## Установка зависимостей
|
||||
|
||||
Для работы скриптов требуются следующие зависимости:
|
||||
|
||||
```bash
|
||||
pip install pymongo
|
||||
```
|
||||
|
||||
#### Codegen
|
||||
```json
|
||||
{
|
||||
"prompt": "Write a Python function that calculates the factorial of a number using recursion.",
|
||||
"expected": "def factorial(n):\n if n == 0 or n == 1:\n return 1\n else:\n return n * factorial(n-1)"
|
||||
}
|
||||
```
|
||||
|
||||
### Валидация
|
||||
|
||||
Скрипт автоматически валидирует generated тесты:
|
||||
- Проверяет наличие обязательных полей (`prompt`, `expected`)
|
||||
- Проверяет, что значения являются строками
|
||||
- Проверяет, что строки не пустые
|
||||
- Поддерживает ручную валидацию существующих тестов через `--validate`
|
||||
|
||||
### Технические детали
|
||||
|
||||
- Скрипт использует ollama_client.py для подключения к Ollama серверу
|
||||
- Каждый generated тест получает уникальный номер (test1.json, test2.json, и т.д.)
|
||||
- Если тест с таким номером уже существует, используется следующий доступный номер
|
||||
- Все тесты сохраняются в формате JSON с UTF-8 кодировкой
|
||||
- Поддерживается любая модель, доступная в Ollama
|
||||
Все зависимости указаны в файле `requirements.txt` в корне проекта.
|
||||
|
||||
73
scripts/convert_json_to_txt.py
Normal file
73
scripts/convert_json_to_txt.py
Normal file
@@ -0,0 +1,73 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Скрипт для конвертации существующих JSON тестов в новый TXT формат.
|
||||
|
||||
Конвертирует все тесты из JSON в TXT с разделителем ================
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
# Добавляем путь к исходникам, чтобы импортировать base
|
||||
sys.path.insert(0, str(Path(__file__).parent.parent))
|
||||
|
||||
from src.benchmarks.base import TEST_SEPARATOR
|
||||
|
||||
def convert_tests(test_dir: str) -> None:
|
||||
"""Конвертирует все тесты в указанной директории."""
|
||||
test_dir_path = Path(test_dir)
|
||||
if not test_dir_path.exists():
|
||||
print(f"❌ Директория {test_dir} не существует")
|
||||
return
|
||||
|
||||
converted_count = 0
|
||||
|
||||
for json_file in test_dir_path.glob("*.json"):
|
||||
try:
|
||||
# Читаем JSON файл
|
||||
with open(json_file, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
|
||||
# Создаем имя TXT файла
|
||||
txt_file = test_dir_path / f"{json_file.stem}.txt"
|
||||
|
||||
# Сохраняем в TXT формате
|
||||
with open(txt_file, "w", encoding="utf-8") as f:
|
||||
f.write(f"{data['prompt']}{TEST_SEPARATOR}{data['expected']}")
|
||||
|
||||
converted_count += 1
|
||||
print(f"✅ Конвертирован {json_file.name} в {txt_file.name}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Ошибка при конвертации {json_file.name}: {str(e)}")
|
||||
|
||||
print(f"\nРезультаты конвертации:")
|
||||
print(f"Конвертировано тестов: {converted_count}")
|
||||
|
||||
def main():
|
||||
"""Основная функция скрипта."""
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Конвертер тестовых данных из JSON в TXT формат",
|
||||
epilog="Примеры использования:\n"
|
||||
" python scripts/convert_json_to_txt.py tests/translation\n"
|
||||
" python scripts/convert_json_to_txt.py tests/summarization\n"
|
||||
" python scripts/convert_json_to_txt.py tests/codegen"
|
||||
)
|
||||
parser.add_argument(
|
||||
"test_dir",
|
||||
type=str,
|
||||
help="Директория с тестами для конвертации (например: tests/translation)"
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
print(f"🔄 Конвертирую тесты в {args.test_dir}")
|
||||
convert_tests(args.test_dir)
|
||||
print("\n✨ Готово!")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
192
scripts/generate_summarization_from_mongo.py
Normal file
192
scripts/generate_summarization_from_mongo.py
Normal file
@@ -0,0 +1,192 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Скрипт для генерации тестов пересказов из MongoDB.
|
||||
|
||||
Извлекает текст статьи из коллекции rssNotification (поле .meta.topicContent)
|
||||
и генерирует тестовые данные в формате JSON для бенчмарка AI.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Dict, Optional
|
||||
|
||||
import pymongo
|
||||
from pymongo import MongoClient
|
||||
|
||||
def connect_to_mongo() -> MongoClient:
|
||||
"""Подключается к MongoDB кластеру."""
|
||||
client = MongoClient(
|
||||
"mongodb://10.0.0.3:27017,10.0.0.4:27017,10.0.0.5:27017/",
|
||||
connectTimeoutMS=30000,
|
||||
socketTimeoutMS=30000,
|
||||
serverSelectionTimeoutMS=30000,
|
||||
retryWrites=True,
|
||||
retryReads=True
|
||||
)
|
||||
return client
|
||||
|
||||
def extract_text_from_topic_content(topic_content: Dict) -> Optional[str]:
|
||||
"""
|
||||
Извлекает текст статьи из .meta.topicContent.
|
||||
|
||||
Args:
|
||||
topic_content: Содержимое поля .meta.topicContent из MongoDB
|
||||
|
||||
Returns:
|
||||
Текст статьи или None, если не удалось извлечь
|
||||
"""
|
||||
if not topic_content:
|
||||
return None
|
||||
|
||||
# Преобразуем в строку, если это не строка
|
||||
content_str = str(topic_content)
|
||||
|
||||
return content_str
|
||||
|
||||
def generate_test_from_mongo_record(record_id: str) -> bool:
|
||||
"""
|
||||
Генерирует тест пересказа из записи MongoDB.
|
||||
|
||||
Args:
|
||||
record_id: ID записи в MongoDB
|
||||
|
||||
Returns:
|
||||
True, если тест успешно generated, False в случае ошибки
|
||||
"""
|
||||
try:
|
||||
client = connect_to_mongo()
|
||||
db = client['tracker_conbot']
|
||||
collection = db['rssNotification']
|
||||
|
||||
# Извлекаем запись по ID
|
||||
record = collection.find_one({"_id": record_id})
|
||||
if not record:
|
||||
print(f"❌ Запись с ID {record_id} не найдена в коллекции")
|
||||
return False
|
||||
|
||||
# Отладочная информация
|
||||
print(f"🔍 Найдена запись: {record_id}")
|
||||
print(f"📋 Полная структура записи:")
|
||||
print(json.dumps(record, ensure_ascii=False, indent=2, default=str))
|
||||
|
||||
# Извлекаем текст из meta.topicContent
|
||||
meta_data = record.get('meta', {})
|
||||
topic_content = meta_data.get('topicContent')
|
||||
if not topic_content:
|
||||
print(f"❌ В записи {record_id} отсутствует поле meta.topicContent")
|
||||
return False
|
||||
|
||||
print(f"📝 Тип поля meta.topicContent: {type(topic_content)}")
|
||||
print(f"📝 Содержимое meta.topicContent (первые 500 символов):")
|
||||
print(str(topic_content)[:500])
|
||||
|
||||
# Извлекаем текст
|
||||
article_text = extract_text_from_topic_content(topic_content)
|
||||
if not article_text:
|
||||
print(f"❌ Не удалось извлечь текст из meta.topicContent записи {record_id}")
|
||||
return False
|
||||
|
||||
print(f"📝 Итоговый текст (первые 500 символов): {article_text[:500]}")
|
||||
|
||||
# Формируем тест
|
||||
test_data = {
|
||||
"prompt": f"Summarize the following text in 1-2 sentences: '{article_text}'",
|
||||
"expected": "" # Ожидаемый результат будет пустым, так как его нужно будет сгенерировать отдельно
|
||||
}
|
||||
|
||||
# Создаем директорию для сохранения теста (всегда в tests/summarization)
|
||||
output_path = Path("tests") / "summarization"
|
||||
output_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Находим следующий доступный номер теста
|
||||
test_num = 1
|
||||
while True:
|
||||
test_file = output_path / f"test{test_num}.json"
|
||||
if not test_file.exists():
|
||||
break
|
||||
test_num += 1
|
||||
|
||||
# Сохраняем тест
|
||||
with open(test_file, "w", encoding="utf-8") as f:
|
||||
json.dump(test_data, f, ensure_ascii=False, indent=2)
|
||||
|
||||
print(f"✅ Создан тест tests/summarization/test{test_num}.json")
|
||||
print(f" Источник: MongoDB запись {record_id}")
|
||||
print(f" Текст статьи (первые 100 символов): {article_text[:100]}...")
|
||||
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Ошибка при генерации теста: {e}")
|
||||
return False
|
||||
finally:
|
||||
if 'client' in locals():
|
||||
client.close()
|
||||
|
||||
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
|
||||
|
||||
return True
|
||||
|
||||
def main():
|
||||
"""Основная функция скрипта."""
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Генератор тестов пересказов из MongoDB",
|
||||
epilog="Примеры использования:\n"
|
||||
" python scripts/generate_summarization_from_mongo.py --record-id 507f1f77bcf86cd799439011 --output-dir tests\n"
|
||||
" python scripts/generate_summarization_from_mongo.py --record-id 507f1f77bcf86cd799439011 --output-dir results"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--record-id",
|
||||
type=str,
|
||||
required=True,
|
||||
help="ID записи в MongoDB (обязательный параметр)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
type=str,
|
||||
default="tests",
|
||||
help="Директория для сохранения generated тестов (по умолчанию: tests)"
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
print(f"🔍 Подключаюсь к MongoDB кластеру...")
|
||||
print(f"📄 Извлекаю запись с ID: {args.record_id}")
|
||||
print(f"💾 Сохраняю тест в: tests/summarization/")
|
||||
|
||||
success = generate_test_from_mongo_record(args.record_id)
|
||||
|
||||
if success:
|
||||
print("\n✨ Готово! Тест успешно generated.")
|
||||
else:
|
||||
print("\n❌ Не удалось generated тест.")
|
||||
sys.exit(1)
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -21,6 +21,7 @@ from typing import Dict, List, Optional
|
||||
sys.path.insert(0, str(Path(__file__).parent.parent))
|
||||
|
||||
from src.models.ollama_client import OllamaClient
|
||||
from src.constants import TEST_SEPARATOR
|
||||
|
||||
def generate_translation_test(ollama: OllamaClient, model: str) -> Dict[str, str]:
|
||||
"""Генерирует тест для перевода через LLM."""
|
||||
@@ -126,23 +127,33 @@ def validate_all_tests(test_dir: str) -> None:
|
||||
valid_count = 0
|
||||
invalid_count = 0
|
||||
|
||||
for json_file in test_dir_path.glob("*.json"):
|
||||
for txt_file in test_dir_path.glob("*.txt"):
|
||||
try:
|
||||
with open(json_file, "r", encoding="utf-8") as f:
|
||||
test_data = json.load(f)
|
||||
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"✅ {json_file.name} - валиден")
|
||||
print(f"✅ {txt_file.name} - валиден")
|
||||
else:
|
||||
invalid_count += 1
|
||||
print(f"❌ {json_file.name} - не валиден")
|
||||
except json.JSONDecodeError:
|
||||
invalid_count += 1
|
||||
print(f"❌ {json_file.name} - некорректный JSON")
|
||||
print(f"❌ {txt_file.name} - не валиден")
|
||||
except Exception as e:
|
||||
invalid_count += 1
|
||||
print(f"❌ {json_file.name} - ошибка: {str(e)}")
|
||||
print(f"❌ {txt_file.name} - ошибка: {str(e)}")
|
||||
|
||||
print(f"\nРезультаты валидации:")
|
||||
print(f"Валидных тестов: {valid_count}")
|
||||
@@ -165,12 +176,12 @@ def generate_tests(ollama: OllamaClient, model: str, count: int, category: str,
|
||||
# Проверяем, существует ли уже тест с таким номером
|
||||
test_num = 1
|
||||
while True:
|
||||
test_file = cat_dir / f"test{test_num}.json"
|
||||
test_file = cat_dir / f"test{test_num}.txt"
|
||||
if not test_file.exists():
|
||||
break
|
||||
test_num += 1
|
||||
|
||||
print(f"🤖 Генерирую тест {cat}/test{test_num}.json...")
|
||||
print(f"🤖 Генерирую тест {cat}/test{test_num}.txt...")
|
||||
|
||||
# Генерируем тест через LLM
|
||||
test_data = generate_test(ollama, model, cat)
|
||||
@@ -180,11 +191,11 @@ def generate_tests(ollama: OllamaClient, model: str, count: int, category: str,
|
||||
print(f"❌ Сгенерирован невалидный тест для {cat}, тест номер {test_num}")
|
||||
continue
|
||||
|
||||
# Сохраняем тест
|
||||
# Сохраняем тест в формате TXT с разделителем
|
||||
with open(test_file, "w", encoding="utf-8") as f:
|
||||
json.dump(test_data, f, ensure_ascii=False, indent=2)
|
||||
f.write(f"{test_data['prompt']}{TEST_SEPARATOR}{test_data['expected']}")
|
||||
|
||||
print(f"✅ Создан тест {cat}/test{test_num}.json")
|
||||
print(f"✅ Создан тест {cat}/test{test_num}.txt")
|
||||
|
||||
def main():
|
||||
"""Основная функция скрипта."""
|
||||
@@ -211,14 +222,12 @@ def main():
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Название модели для генерации тестов (обязательный параметр)"
|
||||
help="Название модели для генерации тестов (требуется при генерации тестов)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ollama-url",
|
||||
type=str,
|
||||
required=True,
|
||||
help="URL подключения к Ollama серверу (обязательный параметр)"
|
||||
help="URL подключения к Ollama серверу (требуется при генерации тестов)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--validate",
|
||||
@@ -231,17 +240,28 @@ def main():
|
||||
if args.validate:
|
||||
print(f"🔍 Начинаю валидацию тестов в {args.validate}")
|
||||
validate_all_tests(args.validate)
|
||||
else:
|
||||
print(f"🤖 Подключаюсь к Ollama серверу: {args.ollama_url}")
|
||||
print(f"📝 Генерирую {args.count} тест(ов) для категории: {args.category}")
|
||||
print(f"🎯 Используемая модель: {args.model}")
|
||||
print("\n✨ Готово!")
|
||||
sys.exit(0)
|
||||
|
||||
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)
|
||||
# Проверяем обязательные параметры для генерации тестов
|
||||
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✨ Готово!")
|
||||
|
||||
|
||||
Reference in New Issue
Block a user