src/ml_training.py - master - Chaizee/ZenithCode_Incident-LLM-analytics

src/ml_training.py
139 строк5 KB
Chaizee
fix: fix speed
12 июн 2026, 16:28
12 июн 2026, 16:28b2745cc
from __future__ import annotations

import json
from pathlib import Path

import joblib
import numpy as np
from loguru import logger
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import LabelEncoder

from config import CATEGORIES, FEW_SHOT_EXAMPLES, ML_EMBEDDING_MODEL, MODELS_DIR, ML_USE_XGBOOST, USE_GPU
from src.device_utils import ml_device
from src.severity import clamp_severity

META_FILE = MODELS_DIR / "model_meta.json"

try:
    from xgboost import XGBClassifier
except ImportError:
    XGBClassifier = None

def _bootstrap_texts() -> tuple[list[str], list[bool], list[int], list[str]]:
    texts, probs, sevs, cats = [], [], [], []
    for ex in FEW_SHOT_EXAMPLES:
        texts.append(ex["text"])
        probs.append(bool(ex.get("is_problem", False)))
        sevs.append(clamp_severity(ex.get("severity", 1)))
        cats.append(str(ex.get("category", "Прочее")))

    extra = [
        ("Слегка шумно от вентиляции в подъезде.", True, 1, "ЖКХ"),
        ("Не убран снег у подъезда, не пройти.", True, 2, "Благоустройство"),
        ("Яма на дворовой дорожке, проехать сложно.", True, 2, "Дороги"),
        ("В доме второй день нет горячей воды.", True, 3, "ЖКХ"),
        ("Прорвало трубу, затопило подвал и квартиры.", True, 4, "ЖКХ"),
        ("Химический выброс, эвакуация жителей.", True, 5, "ЖКХ"),
        ("Благодарю за помощь!", False, 1, "Прочее"),
        ("Во сколько начало соревнований?", False, 1, "Прочее"),
        ("Здравствуйте! В фондах музея есть книга?", False, 1, "Прочее"),
        ("Кто по нему будет ездить?", False, 1, "Прочее"),
    ]
    for t, p, s, c in extra:
        texts.append(t)
        probs.append(p)
        sevs.append(clamp_severity(s))
        cats.append(c)
    return texts, probs, sevs, cats

def save_model_meta(*, embed_dim: int) -> None:
    MODELS_DIR.mkdir(parents=True, exist_ok=True)
    META_FILE.write_text(
        json.dumps(
            {"embedding_model": ML_EMBEDDING_MODEL, "embed_dim": embed_dim},
            ensure_ascii=False,
            indent=2,
        ),
        encoding="utf-8",
    )

def load_model_meta() -> dict | None:
    if not META_FILE.is_file():
        return None
    try:
        return json.loads(META_FILE.read_text(encoding="utf-8"))
    except (json.JSONDecodeError, OSError):
        return None

def models_match_embeddings(embed_dim: int) -> bool:
    if not models_exist():
        return False
    meta = load_model_meta()
    if meta is None:
        return False
    if meta.get("embedding_model") != ML_EMBEDDING_MODEL:
        return False
    if int(meta.get("embed_dim", -1)) != embed_dim:
        return False
    try:
        clf = joblib.load(MODELS_DIR / "problem_clf.joblib")
        return int(getattr(clf, "n_features_in_", -1)) == embed_dim
    except Exception:
        return False

def train_and_save_models(embeddings: np.ndarray) -> None:
    MODELS_DIR.mkdir(parents=True, exist_ok=True)
    texts, probs, sevs, cats = _bootstrap_texts()
    if len(texts) != len(embeddings):
        raise ValueError("Размер bootstrap не совпадает с эмбеддингами")

    embed_dim = int(embeddings.shape[1])
    problem_clf = LogisticRegression(max_iter=500, class_weight="balanced")
    problem_clf.fit(embeddings, probs)
    joblib.dump(problem_clf, MODELS_DIR / "problem_clf.joblib")

    prob_mask = np.array(probs)
    if prob_mask.any():
        Xp = embeddings[prob_mask]
        ys = np.array([clamp_severity(s) for s in sevs])[prob_mask]
        yc = np.array(cats)[prob_mask]
        if ML_USE_XGBOOST and XGBClassifier is not None:
            xgb_kw: dict = {
                "n_estimators": 80,
                "max_depth": 4,
                "learning_rate": 0.1,
            }
            if USE_GPU and ml_device() == "cuda":
                xgb_kw["tree_method"] = "hist"
                xgb_kw["device"] = "cuda"
            sev_enc = LabelEncoder().fit(ys)
            sev_clf = XGBClassifier(
                objective="multi:softmax",
                num_class=len(sev_enc.classes_),
                **xgb_kw,
            )
            sev_clf.fit(Xp, sev_enc.transform(ys))
            cat_enc = LabelEncoder().fit(yc)
            cat_clf = XGBClassifier(
                objective="multi:softmax",
                num_class=len(cat_enc.classes_),
                **xgb_kw,
            )
            cat_clf.fit(Xp, cat_enc.transform(yc))
            joblib.dump(sev_clf, MODELS_DIR / "severity_clf.joblib")
            joblib.dump(sev_enc, MODELS_DIR / "severity_encoder.joblib")
            joblib.dump(cat_clf, MODELS_DIR / "category_clf.joblib")
            joblib.dump(cat_enc, MODELS_DIR / "category_encoder.joblib")
    save_model_meta(embed_dim=embed_dim)
    logger.info("ML-модели сохранены в {} (embed_dim={})", MODELS_DIR, embed_dim)

def models_exist() -> bool:
    needed = [
        "problem_clf.joblib",
        "severity_clf.joblib",
        "severity_encoder.joblib",
        "category_clf.joblib",
        "category_encoder.joblib",
    ]
    return all((MODELS_DIR / f).is_file() for f in needed)