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

src/ml_analysis.py
81 строка2 KB
Chaizee
feat: new realisation
11 июн 2026, 19:58
11 июн 2026, 19:58697b236
from __future__ import annotations

from typing import Any

import polars as pl

from config import (
    CLUSTER_MIN_SIZE,
    CLUSTER_SIMILARITY_MIN,
    MAX_PROBLEMS_FOR_CLUSTERING,
    USE_HYBRID_CLUSTERING,
)
from src.memory_guard import free_memory
from src.ml_classifier import MLClassifierBackend

def analyze_dataframe(
    df: pl.DataFrame,
    backend: MLClassifierBackend,
    *,
    llm: Any = None,
) -> pl.DataFrame:
    texts = df.get_column("incident_text").cast(pl.Utf8).to_list()
    themes = (
        df.get_column("theme").cast(pl.Utf8).to_list()
        if "theme" in df.columns
        else [""] * len(texts)
    )
    rows = backend.classify_batch(texts, themes=themes, llm=llm)
    del texts, themes

    enriched = df.with_columns(
        pl.Series("is_problem", [r["is_problem"] for r in rows]),
        pl.Series("severity", [r["severity"] for r in rows]),
        pl.Series("category", [r["category"] for r in rows]),
    )
    del rows

    n_problems = enriched.filter(pl.col("is_problem")).height
    if (
        USE_HYBRID_CLUSTERING
        and len(enriched) >= CLUSTER_MIN_SIZE
        and n_problems <= MAX_PROBLEMS_FOR_CLUSTERING
    ):
        enriched = _cluster_problems(enriched, backend)
    else:
        enriched = enriched.with_columns(pl.lit(-1).alias("cluster"))
    return enriched

def _cluster_problems(df: pl.DataFrame, backend: MLClassifierBackend) -> pl.DataFrame:
    indexed = df.with_row_index("_idx")
    problems = indexed.filter(pl.col("is_problem"))
    if len(problems) < CLUSTER_MIN_SIZE:
        return df.with_columns(pl.lit(-1).alias("cluster"))

    try:
        from sklearn.cluster import DBSCAN
    except ImportError:
        return df.with_columns(pl.lit(-1).alias("cluster"))

    texts = problems.get_column("incident_text").cast(pl.Utf8).to_list()
    idxs = problems["_idx"].to_list()
    emb = backend.encode(texts)
    del texts
    try:
        labels = DBSCAN(
            eps=1 - CLUSTER_SIMILARITY_MIN,
            min_samples=CLUSTER_MIN_SIZE,
            metric="cosine",
        ).fit_predict(emb)
    finally:
        del emb
        free_memory()

    cluster_by_row = dict(zip(idxs, labels.tolist(), strict=False))
    del idxs, labels

    return indexed.with_columns(
        pl.col("_idx")
        .map_elements(lambda i: cluster_by_row.get(i, -1), return_dtype=pl.Int64)
        .alias("cluster")
    ).drop("_idx")