101 lines
3.0 KiB
Python
101 lines
3.0 KiB
Python
import logging
|
||
import time
|
||
from typing import Dict, Any, List
|
||
from abc import ABC, abstractmethod
|
||
from models.ollama_client import OllamaClient
|
||
|
||
class Benchmark(ABC):
|
||
"""Базовый класс для всех бенчмарков."""
|
||
|
||
def __init__(self, name: str):
|
||
"""
|
||
Инициализация бенчмарка.
|
||
|
||
Args:
|
||
name: Название бенчмарка
|
||
"""
|
||
self.name = name
|
||
self.logger = logging.getLogger(__name__)
|
||
|
||
@abstractmethod
|
||
def load_test_data(self) -> List[Dict[str, Any]]:
|
||
"""
|
||
Загрузка тестовых данных.
|
||
|
||
Returns:
|
||
Список тестовых случаев
|
||
"""
|
||
pass
|
||
|
||
@abstractmethod
|
||
def evaluate(self, model_response: str, expected: str) -> float:
|
||
"""
|
||
Оценка качества ответа модели.
|
||
|
||
Args:
|
||
model_response: Ответ от модели
|
||
expected: Ожидаемый ответ
|
||
|
||
Returns:
|
||
Метрика качества (0-1)
|
||
"""
|
||
pass
|
||
|
||
def run(self, ollama_client: OllamaClient, model_name: str) -> Dict[str, Any]:
|
||
"""
|
||
Запуск бенчмарка.
|
||
|
||
Args:
|
||
ollama_client: Клиент для работы с Ollama
|
||
model_name: Название модели
|
||
|
||
Returns:
|
||
Результаты бенчмарка
|
||
"""
|
||
test_cases = self.load_test_data()
|
||
results = []
|
||
|
||
for i, test_case in enumerate(test_cases, 1):
|
||
try:
|
||
self.logger.info(f"Running test case {i}/{len(test_cases)} for {self.name}")
|
||
|
||
# Замер времени
|
||
start_time = time.time()
|
||
|
||
# Получение ответа от модели
|
||
prompt = test_case['prompt']
|
||
model_response = ollama_client.generate(
|
||
model=model_name,
|
||
prompt=prompt,
|
||
options={'temperature': 0.7}
|
||
)
|
||
|
||
# Замер времени
|
||
latency = time.time() - start_time
|
||
|
||
# Оценка качества
|
||
score = self.evaluate(model_response, test_case['expected'])
|
||
|
||
results.append({
|
||
'test_case': test_case['name'],
|
||
'prompt': prompt,
|
||
'expected': test_case['expected'],
|
||
'model_response': model_response,
|
||
'score': score,
|
||
'latency': latency
|
||
})
|
||
|
||
except Exception as e:
|
||
self.logger.error(f"Error in test case {i}: {e}")
|
||
results.append({
|
||
'test_case': test_case['name'],
|
||
'error': str(e)
|
||
})
|
||
|
||
return {
|
||
'benchmark_name': self.name,
|
||
'total_tests': len(test_cases),
|
||
'successful_tests': len([r for r in results if 'score' in r]),
|
||
'results': results
|
||
}
|