chapter4/main.py - master - dhdied/course_project

chapter4/main.py
146 строк6 KB
dhdied
final version
02 июн 2026, 03:53
02 июн 2026, 03:53080b2bc
import pandas as pd
import json
import os
import matplotlib.pyplot as plt

base_dir = os.getcwd()
merged_db_dir = os.path.join(base_dir, 'merged_database')

excel_files = [
    'binary.xlsx',
    'dates.xlsx',
    'dev_to_check.xlsx',
    'english_for_remark_NONE&DBP.xlsx',
    'fuzzy.xlsx',
    'Список запросов_db_check_v2.xlsx'
]

SYSTEM_INSTRUCTION = "Используя предоставленную схему базы данных, напиши синтаксически корректный SQL-запрос для ответа на вопрос пользователя."

def get_clean_schema(db_id):
    db_id = str(db_id).strip()
    schema_path = os.path.join(merged_db_dir, db_id, 'schema.sql')
    
    if os.path.exists(schema_path):
        with open(schema_path, 'r', encoding='utf-8') as f:
            lines = f.readlines()
        
        clean_lines = []
        for line in lines:
            line_upper = line.strip().upper()
            # ФИЛЬТР: Убираем дамп данных и системные команды, оставляем только структуру
            if line_upper.startswith(('INSERT', 'PRAGMA', 'BEGIN', 'COMMIT', '/*')):
                continue
            if line.strip(): # Игнорируем пустые строки
                clean_lines.append(line)
                
        return "".join(clean_lines).strip()
    
    return None 

def determine_complexity(sql):
    sql_upper = str(sql).upper()
    if 'INTERSECT' in sql_upper or 'EXCEPT' in sql_upper or 'UNION' in sql_upper:
        return 'Экстра-сложный'
    elif sql_upper.count('JOIN') > 1 or 'HAVING' in sql_upper or 'LIMIT' in sql_upper:
        return 'Сложный'
    elif 'JOIN' in sql_upper or 'GROUP BY' in sql_upper:
        return 'Средний'
    else:
        return 'Базовый'

dataset = []
question_lengths = []
sql_lengths = []
sql_complexities = {'Базовый': 0, 'Средний': 0, 'Сложный': 0, 'Экстра-сложный': 0}

skipped_no_schema = 0

print("Начинаем умную очистку и сборку датасета...\n")

for file in excel_files:
    file_path = os.path.join(base_dir, file)
    if not os.path.exists(file_path):
        continue

    df = pd.read_excel(file_path)
    df.rename(columns=lambda x: str(x).strip().lower(), inplace=True)
    
    q_priorities = ['ru_corrected', 'ru', 'question', 'en']
    sql_priorities = ['sql_ru_corrected', 'sql_ru', 'sql_en_corrected', 'sql_en', 'ru_sql', 'en_sql', 'sql', 'query']
    
    q_col = next((col for col in q_priorities if col in df.columns), None)
    sql_col = next((col for col in sql_priorities if col in df.columns), None)
    db_col = 'db_id' if 'db_id' in df.columns else None

    if not q_col or not sql_col or not db_col:
        continue

    print(f"Читаем {file} (строк: {len(df)})")

    for index, row in df.iterrows():
        if pd.isna(row[q_col]) or pd.isna(row[sql_col]) or pd.isna(row[db_col]):
            continue

        db_id = str(row[db_col]).strip()
        question = str(row[q_col]).strip()
        sql = str(row[sql_col]).strip()

        if len(question) < 5 or len(sql) < 5:
            continue

        schema = get_clean_schema(db_id)
        
        # ФИЛЬТР: Пропускаем строки, для которых нет схемы БД
        if not schema:
            skipped_no_schema += 1
            continue

        input_text = f"Схема базы данных:\n{schema}\n\nВопрос: {question}"

        dataset.append({
            "instruction": SYSTEM_INSTRUCTION,
            "input": input_text,
            "output": sql
        })

        question_lengths.append(len(question.split()))
        sql_lengths.append(len(sql.split()))

        complexity = determine_complexity(sql)
        sql_complexities[complexity] += 1

output_dir = os.path.join(base_dir, 'data')
os.makedirs(output_dir, exist_ok=True)
output_file = os.path.join(output_dir, 'train.jsonl')

with open(output_file, 'w', encoding='utf-8') as f:
    for item in dataset:
        f.write(json.dumps(item, ensure_ascii=False) + '\n')

print(f"\n Файл сохранен в: {output_file}")
print(f"Обраработано чистых пар: {len(dataset)}")
print(f"Отброшено строк (нет схемы БД): {skipped_no_schema}")

print("\n--- СТАТИСТИКА ДЛЯ ПАСПОРТА ДАТАСЕТА ---")
avg_q_len = sum(question_lengths) / len(question_lengths) if question_lengths else 0
avg_sql_len = sum(sql_lengths) / len(sql_lengths) if sql_lengths else 0

print(f"Средняя длина вопроса: {avg_q_len:.2f}")
print(f"Средняя длина SQL-запроса: {avg_sql_len:.2f}")
print(f"Распределение: {sql_complexities}")

plt.figure(figsize=(8, 5))
bars = plt.bar(sql_complexities.keys(), sql_complexities.values(), color=['#4CAF50', '#2196F3', '#FF9800', '#F44336'])
plt.title('Распределение классов (сложность SQL-запросов)', fontsize=14)
plt.xlabel('Уровень сложности', fontsize=12)
plt.ylabel('Количество примеров', fontsize=12)
plt.grid(axis='y', linestyle='--', alpha=0.7)

for bar in bars:
    yval = bar.get_height()
    plt.text(bar.get_x() + bar.get_width()/2, yval + (max(sql_complexities.values())*0.01),
             int(yval), ha='center', va='bottom', fontsize=11)

plt.tight_layout()
plt.show()