src/train.py - ui - Chaizee/ZenithCode_Incident-LLM-analytics

src/train.py
213 строк8 KB
Chaizee
feat: add llm-analyzer realisation
08 июн 2026, 18:44
08 июн 2026, 18:440796e1d
from __future__ import annotations

import json
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any, Callable

import pandas as pd
from loguru import logger

from config import (
    BASE_XLSX,
    FEW_SHOT_COUNT,
    FEW_SHOT_FILE,
    LABELED_BASE,
    OUTPUT_DIR,
    TEST_METRICS,
    TEST_XLSX,
    TRAIN_METRICS,
    TRAINING_DIR,
)
from src.analytics import build_analytics
from src.data import iter_chunks, load_excel, merge_llm
from src.llm import LLMEngine
from src.pipeline import PipelineResult, run_pipeline
from src.report import build_excel, save_excel

ProgressFn = Callable[[float, str], None]


@dataclass
class TrainResult:
    rows_labeled: int = 0
    problems_found: int = 0
    few_shot_count: int = 0
    metrics: dict[str, Any] = field(default_factory=dict)
    logs: list[str] = field(default_factory=list)


def _noop(p: float, m: str) -> None:
    pass


def _dataset_metrics(df: pd.DataFrame) -> dict[str, Any]:
    problems = df[df["is_problem"].astype(bool)] if "is_problem" in df.columns else df.iloc[0:0]
    return {
        "total_rows": len(df),
        "problems": len(problems),
        "problem_rate_pct": round(100 * len(problems) / max(len(df), 1), 2),
        "mean_severity": round(float(problems["severity"].mean()), 2) if len(problems) else 0,
        "municipalities": int(df["municipality"].nunique()) if "municipality" in df.columns else 0,
        "top_issues": problems["core_issue"].value_counts().head(10).to_dict() if len(problems) else {},
    }


def _build_few_shot(labeled: pd.DataFrame, n: int = FEW_SHOT_COUNT) -> list[dict[str, Any]]:
    """Примеры для подсказок LLM: проблемы и не-проблемы."""
    examples: list[dict[str, Any]] = []
    if labeled.empty or "is_problem" not in labeled.columns:
        return examples

    half = max(1, n // 2)
    for flag in (True, False):
        sub = labeled[labeled["is_problem"].astype(bool) == flag]
        sub = sub[sub["incident_text"].astype(str).str.len() > 20]
        if sub.empty:
            continue
        sample = sub.sample(min(half, len(sub)), random_state=42)
        for _, row in sample.iterrows():
            examples.append(
                {
                    "text": str(row["incident_text"])[:500],
                    "is_problem": bool(row["is_problem"]),
                    "severity": int(row["severity"]),
                    "core_issue": str(row["core_issue"]),
                }
            )
    return examples[:n]


def train_on_base(
    *,
    engine: LLMEngine | None = None,
    source: Path | None = None,
    max_rows: int | None = None,
    skip_llm: bool = False,
    progress: ProgressFn | None = None,
) -> TrainResult:
    """
    Обучение = прогон LLM по base.xlsx → размеченный parquet + SQLite-кэш + few-shot.
    Полный base.xlsx (~419k) на GPU займёт много часов; для отладки задайте max_rows.
    """
    report = progress or _noop
    result = TrainResult()
    path = source or BASE_XLSX

    if not path.is_file():
        raise FileNotFoundError(f"Не найден файл обучения: {path}")

    report(0.02, f"Чтение {path.name}…")
    df = load_excel(path, nrows=max_rows)
    result.logs.append(f"Обучение: {len(df):,} строк из {path.name}")
    logger.info("Train dataset: {} rows", len(df))

    engine = engine or LLMEngine(use_cache=True)
    if not skip_llm:
        report(0.05, "Загрузка LLM…")
        try:
            engine.load()
            result.logs.append(f"Модель: {engine.model_name}")
        except Exception as exc:
            logger.warning("LLM недоступна: {}", exc)
            result.logs.append(f"LLM fallback: {exc}")

    frames: list[pd.DataFrame] = []
    total = len(df)
    processed = 0

    for i, chunk in enumerate(iter_chunks(df), start=1):
        texts = chunk["incident_text"].astype(str).tolist()
        if skip_llm:
            llm_rows = [{"is_problem": False, "severity": 1, "core_issue": "—"} for _ in texts]
        else:

            def batch_prog(done: int, tot: int) -> None:
                frac = (processed + done) / max(total, 1)
                report(0.1 + 0.7 * frac, f"Разметка {processed + done:,}/{total:,}")

            llm_rows = engine.analyze_all(texts, on_progress=batch_prog)

        frames.append(merge_llm(chunk, llm_rows))
        processed += len(chunk)
        report(0.1 + 0.7 * (processed / max(total, 1)), f"Чанк {i}")

    labeled = pd.concat(frames, ignore_index=True)
    TRAINING_DIR.mkdir(parents=True, exist_ok=True)
    labeled.to_parquet(LABELED_BASE, index=False)

    few_shot = _build_few_shot(labeled)
    FEW_SHOT_FILE.write_text(json.dumps(few_shot, ensure_ascii=False, indent=2), encoding="utf-8")
    engine.set_few_shot(few_shot)

    metrics = _dataset_metrics(labeled)
    metrics["trained_at"] = datetime.now().isoformat()
    metrics["source"] = str(path)
    metrics["max_rows"] = max_rows
    TRAIN_METRICS.write_text(json.dumps(metrics, ensure_ascii=False, indent=2), encoding="utf-8")

    result.rows_labeled = len(labeled)
    result.problems_found = metrics["problems"]
    result.few_shot_count = len(few_shot)
    result.metrics = metrics
    result.logs.append(f"Сохранено: {LABELED_BASE.name}, few-shot: {len(few_shot)}")
    report(1.0, "Обучение завершено")
    logger.info("Train done: {} rows, {} problems", len(labeled), metrics["problems"])
    return result


def test_on_test(
    *,
    engine: LLMEngine | None = None,
    source: Path | None = None,
    use_few_shot: bool = True,
    skip_llm: bool = False,
    progress: ProgressFn | None = None,
) -> PipelineResult:
    """Тест на test.xlsx с отчётом и метриками (использует кэш после обучения)."""
    path = source or TEST_XLSX
    if not path.is_file():
        raise FileNotFoundError(f"Не найден тестовый файл: {path}")

    engine = engine or LLMEngine(use_cache=True)
    if use_few_shot and FEW_SHOT_FILE.is_file():
        try:
            few_shot = json.loads(FEW_SHOT_FILE.read_text(encoding="utf-8"))
            engine.set_few_shot(few_shot)
        except Exception as exc:
            logger.warning("Few-shot не загружен: {}", exc)

    result = run_pipeline(
        path,
        engine=engine,
        skip_llm=skip_llm,
        region_name="Омская область (тест)",
        progress=progress,
    )

    test_metrics = _dataset_metrics(result.enriched)
    test_metrics["tested_at"] = datetime.now().isoformat()
    test_metrics["source"] = str(path)
    if TRAIN_METRICS.is_file():
        try:
            train_m = json.loads(TRAIN_METRICS.read_text(encoding="utf-8"))
            test_metrics["train_problem_rate_pct"] = train_m.get("problem_rate_pct")
            test_metrics["delta_problem_rate"] = round(
                test_metrics["problem_rate_pct"] - float(train_m.get("problem_rate_pct", 0)), 2
            )
        except Exception:
            pass

    TRAINING_DIR.mkdir(parents=True, exist_ok=True)
    TEST_METRICS.write_text(json.dumps(test_metrics, ensure_ascii=False, indent=2), encoding="utf-8")
    result.logs.append(f"Метрики теста: {TEST_METRICS.name}")

    stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    save_excel(
        result.analytics["top10"],
        result.analytics["top3"],
        OUTPUT_DIR / f"test_report_{stamp}.xlsx",
    )
    return result