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

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

from typing import Any

import joblib
import numpy as np
from loguru import logger

from config import (
    CATEGORIES,
    ML_EMBEDDING_MODEL,
    ML_ENCODE_MAX_CHARS,
    ML_GRAY_LLM_ENABLED,
    ML_GRAY_LLM_MAX_SAMPLES,
    ML_GRAY_ZONE_LOW,
    ML_PROBLEM_THRESHOLD,
    MODELS_DIR,
)
from src.device_utils import configure_cpu_threads, ml_classify_chunk, ml_device, ml_embed_batch_size
from src.memory_guard import free_memory
from src.ml_training import models_match_embeddings, train_and_save_models
from src.problem_heuristics import adjust_classification
from src.severity import apply_problem_flags, clamp_severity, refine_severity

class MLClassifierBackend:
    def __init__(self) -> None:
        self._encoder: Any = None
        self._problem_clf: Any = None
        self._severity_clf: Any = None
        self._severity_enc: Any = None
        self._category_clf: Any = None
        self._category_enc: Any = None

    @property
    def ready(self) -> bool:
        return self._encoder is not None and self._problem_clf is not None

    def load(self) -> None:
        try:
            from sentence_transformers import SentenceTransformer
        except ImportError as exc:
            raise ImportError(
                "Пакет sentence-transformers не установлен. "
                "Выполните: pip install -r requirements.txt"
            ) from exc

        device = ml_device()
        if device == "cpu":
            threads = configure_cpu_threads()
            logger.info("PyTorch CPU threads: {}", threads)
        logger.info("Загрузка эмбеддингов {} на {}", ML_EMBEDDING_MODEL, device)
        self._encoder = SentenceTransformer(ML_EMBEDDING_MODEL, device=device)

        from src.ml_training import _bootstrap_texts

        probe = self.encode(["probe"])
        embed_dim = int(probe.shape[1])
        del probe

        if not models_match_embeddings(embed_dim):
            logger.warning(
                "ML-модели несовместимы с {} (dim={}) — переобучение bootstrap…",
                ML_EMBEDDING_MODEL,
                embed_dim,
            )
            texts, _, _, _ = _bootstrap_texts()
            emb = self.encode(texts)
            train_and_save_models(emb)
            del emb, texts

        self._problem_clf = joblib.load(MODELS_DIR / "problem_clf.joblib")
        self._severity_clf = joblib.load(MODELS_DIR / "severity_clf.joblib")
        sev_enc_path = MODELS_DIR / "severity_encoder.joblib"
        self._severity_enc = joblib.load(sev_enc_path) if sev_enc_path.is_file() else None
        self._category_clf = joblib.load(MODELS_DIR / "category_clf.joblib")
        self._category_enc = joblib.load(MODELS_DIR / "category_encoder.joblib")
        logger.info("ML-классификаторы загружены")

    def unload_encoder(self) -> None:
        if self._encoder is not None:
            del self._encoder
            self._encoder = None
            free_memory(deep=True)
            logger.info("Эмбеддинг-модель выгружена")

    def unload(self) -> None:
        self.unload_encoder()

    def dispose(self) -> None:
        self.unload_encoder()
        for attr in ("_problem_clf", "_severity_clf", "_severity_enc", "_category_clf", "_category_enc"):
            obj = getattr(self, attr, None)
            if obj is not None:
                del obj
                setattr(self, attr, None)
        free_memory(deep=True)
        logger.info("ML-бэкенд полностью выгружен")

    @staticmethod
    def _trim_texts(texts: list[str]) -> list[str]:
        limit = max(64, int(ML_ENCODE_MAX_CHARS))
        return [str(t or "")[:limit] for t in texts]

    def encode(self, texts: list[str]) -> np.ndarray:
        if self._encoder is None:
            raise RuntimeError("ML backend не загружен")
        trimmed = self._trim_texts(texts)
        return np.asarray(
            self._encoder.encode(
                trimmed,
                batch_size=ml_embed_batch_size(),
                show_progress_bar=False,
                normalize_embeddings=True,
                convert_to_numpy=True,
            ),
            dtype=np.float32,
        )

    @staticmethod
    def _decode_severity(raw_pred: np.ndarray) -> np.ndarray:
        return np.array([clamp_severity(v) for v in raw_pred], dtype=int)

    def _classify_chunk(
        self,
        texts: list[str],
        themes: list[str],
    ) -> tuple[list[dict[str, Any]], list[int]]:
        emb = self.encode(texts)
        try:
            prob_scores = self._problem_clf.predict_proba(emb)[:, 1]
            is_problem = prob_scores >= ML_PROBLEM_THRESHOLD

            gray_idx = [
                i
                for i, p in enumerate(prob_scores)
                if ML_GRAY_ZONE_LOW <= p < ML_PROBLEM_THRESHOLD
            ]

            sev_pred = self._severity_clf.predict(emb)
            if self._severity_enc is not None:
                raw_sev = self._severity_enc.inverse_transform(sev_pred)
            else:
                raw_sev = sev_pred
            sev = self._decode_severity(raw_sev)

            cat_idx = self._category_clf.predict(emb)
            categories = self._category_enc.inverse_transform(cat_idx)

            rows: list[dict[str, Any]] = []
            for i in range(len(texts)):
                prob_flag = bool(is_problem[i])
                theme = themes[i] if i < len(themes) else ""
                sev_val = refine_severity(
                    int(sev[i]),
                    texts[i],
                    theme=theme,
                    is_problem=prob_flag,
                )
                prob, sev_val = apply_problem_flags(prob_flag, sev_val)
                row = adjust_classification(
                    texts[i],
                    {
                        "is_problem": prob,
                        "severity": sev_val,
                        "category": str(categories[i]) if categories[i] in CATEGORIES else "Прочее",
                    },
                )
                rows.append(row)
            return rows, gray_idx
        finally:
            del emb

    def _apply_gray_llm(
        self,
        rows: list[dict[str, Any]],
        texts: list[str],
        themes: list[str],
        gray_indices: list[int],
        llm: Any,
    ) -> None:
        if not ML_GRAY_LLM_ENABLED or not gray_indices or llm is None or not llm.ready:
            return
        from src.llm_tasks import verify_problems_batch

        cap = max(0, int(ML_GRAY_LLM_MAX_SAMPLES))
        if cap and len(gray_indices) > cap:
            logger.info(
                "Серая зона LLM: {} из {} строк (лимит {})",
                cap,
                len(gray_indices),
                cap,
            )
            gray_indices = gray_indices[:cap]

        gray_texts = [texts[i] for i in gray_indices]
        verified = verify_problems_batch(llm, gray_texts)
        for j, i in enumerate(gray_indices):
            if rows[i]["is_problem"]:
                continue
            flag = verified[j]
            if not flag:
                continue
            theme = themes[i] if i < len(themes) else ""
            sev_val = refine_severity(
                int(rows[i]["severity"]),
                texts[i],
                theme=theme,
                is_problem=flag,
            )
            prob, sev_val = apply_problem_flags(flag, sev_val)
            rows[i]["is_problem"] = prob
            rows[i]["severity"] = sev_val

    def classify_batch(
        self,
        texts: list[str],
        *,
        themes: list[str] | None = None,
        llm: Any = None,
    ) -> list[dict[str, Any]]:
        if not texts:
            return []

        themes = themes or [""] * len(texts)
        out: list[dict[str, Any]] = []
        gray_global: list[int] = []

        chunk_size = ml_classify_chunk()
        for start in range(0, len(texts), chunk_size):
            chunk = texts[start : start + chunk_size]
            chunk_themes = themes[start : start + chunk_size]
            chunk_rows, gray_idx = self._classify_chunk(chunk, chunk_themes)
            gray_global.extend(start + i for i in gray_idx)
            out.extend(chunk_rows)
            if start > 0 and start % (chunk_size * 2) == 0:
                free_memory()

        self._apply_gray_llm(out, texts, themes, gray_global, llm)
        return out