run_benchmarks.py - master - sonne/Task-RatsAD-31.group

run_benchmarks.py
131 строка6 KB
arseon21
Add SQL
24 июн 2025, 20:24
24 июн 2025, 20:244416e94
import psycopg2
import time
import uuid
import pandas as pd
from psycopg2 import extras

# НАСТРОЙКИ ПОДКЛЮЧЕНИЯ К БАЗЕ ДАННЫХ

DB_PARAMS = {
    'dbname': 'fuzzy_search_lab',  # Имя базы данных
    'user': 'postgres',           # Имя пользователя
    'password': '1963',             # Ваш пароль PostgreSQL
    'host': 'localhost',
    'port': '5432'
}


class FuzzySearchBenchmark:
    def __init__(self, connection_params):
        self.conn = psycopg2.connect(**connection_params)
        self.test_run_id = str(uuid.uuid4())
        print(f"Новый запуск тестов. ID: {self.test_run_id}")

    def _calculate_precision_recall(self, found_items: set, relevant_items: set) -> tuple[float, float]:
        if not relevant_items and not found_items: return 1.0, 1.0
        if not relevant_items or not found_items: return 0.0, 0.0
        true_positives = len(found_items.intersection(relevant_items))
        precision = true_positives / len(found_items) if found_items else 0
        recall = true_positives / len(relevant_items) if relevant_items else 0
        return precision, recall

    def _f1_score(self, precision: float, recall: float) -> float:
        if precision + recall == 0: return 0.0
        return 2 * (precision * recall) / (precision + recall)

    def _log_results(self, method_name, dataset_size, query_text, execution_time_ms, result_count, index_used):
        with self.conn.cursor() as cursor:
            sql = """INSERT INTO search_benchmarks
                     (method, dataset_size, query_text, execution_time_ms, result_count, index_used, test_run_id)
                     VALUES (%s, %s, %s, %s, %s, %s, %s)"""
            cursor.execute(sql, (method_name, dataset_size, query_text, execution_time_ms, result_count, index_used,
                                 self.test_run_id))
            self.conn.commit()

    def benchmark_method(self, query_template: str, search_term: str) -> tuple[float, int, set, bool]:
        query = self.conn.cursor().mogrify(query_template, (search_term,))
        index_used = False
        start_time = time.perf_counter()
        with self.conn.cursor() as cursor:
            cursor.execute(query)
            results = cursor.fetchall()
        end_time = time.perf_counter()
        execution_time_ms = (end_time - start_time) * 1000
        found_items = {row[0] for row in results}
        result_count = len(found_items)
        return execution_time_ms, result_count, found_items, index_used

    def run_all_benchmarks(self):
        print("\nЗапуск бенчмарков...")

        with self.conn.cursor() as cursor:

            test_cases = [
                ('computer', 'copmuter', 'transposition'), ('monitor', 'mointor', 'transposition'),
                ('keyboard', 'keybord', 'deletion'), ('software', 'sofware', 'deletion'),
                ('processor', 'processsor', 'insertion'), ('adapter', 'adappter', 'insertion'),
                ('mouse', 'mouce', 'substitution'), ('windows', 'windovs', 'substitution')
            ]
            cursor.execute("SELECT COUNT(*) FROM products;")
            dataset_size = cursor.fetchone()[0]

        query_templates = {
            'ILIKE': "SELECT name FROM products WHERE name ILIKE '%%' || %s || '%%';",
            'Trigram (%)': "SELECT name FROM products WHERE name %% %s;",
            'Levenshtein (<=3)': "SELECT name FROM products WHERE levenshtein(name, %s) <= 3;",
            'Soundex': "SELECT name FROM products WHERE soundex(name) = soundex(%s);",
            'FTS (@@)': "SELECT name FROM products WHERE to_tsvector('english', name) @@ plainto_tsquery('english', %s);"
        }

        for correct_term, typo_term, error_type in test_cases:
            print(f"\nТестирование ошибки '{error_type}' (Правильно: '{correct_term}', Ищем: '{typo_term}')")

            with self.conn.cursor() as cursor:

                cursor.execute("SELECT name FROM products WHERE name ILIKE '%%' || %s || '%%'", (correct_term,))
                relevant_items = {row[0] for row in cursor.fetchall()}
                if not relevant_items:
                    print(f"  НЕ НАЙДЕНО эталонных записей для '{correct_term}'. Пропускаем оценку точности.")
                    continue

            for method, template in query_templates.items():
                exec_time, res_count, found_items, idx_used = self.benchmark_method(template, typo_term)
                precision, recall = self._calculate_precision_recall(found_items, relevant_items)
                f1 = self._f1_score(precision, recall)
                self._log_results(method, dataset_size, typo_term, exec_time, res_count, idx_used)
                print(
                    f"  Метод: {method:<20} | "
                    f"Время: {exec_time:<7.2f} мс | "
                    f"Найдено: {res_count:<3} | "
                    f"Точность: {precision:.2f} | "
                    f"Полнота: {recall:.2f} | "
                    f"F1: {f1:.2f} | "
                    f"Индекс: {idx_used}"
                )

    def close(self):
        if self.conn:
            self.conn.close()


if __name__ == '__main__':

    print("Подготовка FTS колонки...")
    try:
        conn = psycopg2.connect(**DB_PARAMS)
        with conn.cursor() as cur:
            cur.execute("ALTER TABLE products ADD COLUMN IF NOT EXISTS search_vector tsvector;")
            cur.execute("UPDATE products SET search_vector = to_tsvector('english', name || ' ' || description);")
            conn.commit()
    except Exception as e:
        print(f"Ошибка подготовки FTS: {e}")
    finally:
        if 'conn' in locals(): conn.close()

    benchmark_runner = FuzzySearchBenchmark(DB_PARAMS)
    try:
        benchmark_runner.run_all_benchmarks()
    except Exception as e:
        print(f"Произошла ошибка во время выполнения бенчмарков: {e}")
    finally:
        benchmark_runner.close()