Update app.py

app.py

@@ -4,37 +4,60 @@ from html import unescape
 import gradio as gr
 import numpy as np
 
-# ====== TF-IDF + LR (joblib / sklearn) ======
-PIPE = None
-try:
+# ==============================
+# TF-IDF + LR (joblib / sklearn)
+# ==============================
+PIPE, PIPE_PATH = None, None
+def _load_tfidf():
     import joblib
-    PIPE = joblib.load("model.joblib")  # сохранённый пайплайн TF-IDF+LR
+    candidates = [
+        "model.joblib",
+        "artifacts/model.joblib",
+        "tfidf/model.joblib",
+    ]
+    last_err = None
+    for p in candidates:
+        if os.path.exists(p):
+            try:
+                pipe = joblib.load(p)
+                return pipe, p
+            except Exception as e:
+                last_err = e
+    if last_err:
+        print(f"[WARN] TF-IDF load failed: {last_err}")
+    else:
+        print("[WARN] TF-IDF model not found in", candidates)
+    return None, None
+
+try:
+    PIPE, PIPE_PATH = _load_tfidf()
 except Exception as e:
-    PIPE = None
-    print(f"[WARN] Не удалось загрузить model.joblib: {e}")
+    print(f"[WARN] Не удалось инициализировать TF-IDF: {e}")
+    PIPE, PIPE_PATH = None, None
 
-# ====== Transformer (ruBERT-tiny2) ======
+# ==============================
+# Transformer (ruBERT-tiny2)
+# ==============================
 TRANSFORMER = {"model": None, "tokenizer": None, "device": "cpu"}
 try:
     import torch
     from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
-    MODEL_DIR = "."  # в корне лежат config.json, model.safetensors, tokenizer.*
+    MODEL_DIR = "."  # файлы трансформера лежат в корне Space
     device = "cuda" if torch.cuda.is_available() else "cpu"
 
     tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR, local_files_only=True)
     model = AutoModelForSequenceClassification.from_pretrained(MODEL_DIR, local_files_only=True)
     model.to(device).eval()
 
-    TRANSFORMER["model"] = model
-    TRANSFORMER["tokenizer"] = tokenizer
-    TRANSFORMER["device"] = device
+    TRANSFORMER.update({"model": model, "tokenizer": tokenizer, "device": device})
 except Exception as e:
     print(f"[WARN] Не удалось загрузить ruBERT: {e}")
 
-# ====== Порог по умолчанию ======
-DEFAULT_THRESHOLD = 0.70  # как просили
-# если есть inference_config.json от обучения трансформера — подхватим рекомендованный
+# ==============================
+# Порог по умолчанию
+# ==============================
+DEFAULT_THRESHOLD = 0.70
 try:
     if os.path.exists("inference_config.json"):
         with open("inference_config.json", "r", encoding="utf-8") as f:
@@ -43,7 +66,9 @@ try:
 except Exception:
     pass
 
-# ====== Предобработка для трансформера (как в обучении) ======
+# ==============================
+# Предобработка (как при обучении трансформера)
+# ==============================
 from nltk.stem.snowball import RussianStemmer
 stemmer = RussianStemmer(ignore_stopwords=False)
 
@@ -69,29 +94,52 @@ def clean_and_stem(s: str) -> str:
         out.append(t if t in {"url", "tag", "num"} else stemmer.stem(t))
     return " ".join(out)
 
-# ====== Инференс ======
-def infer_tfidf(text: str) -> float:
-    """Вернёт P(toxic) из TF-IDF+LR. В пайплайне уже есть свой preprocessor."""
+# ==============================
+# Инференс
+# ==============================
+def infer_tfidf(text: str):
+    """
+    Вернёт (proba, err_msg). Если всё ок: (float in [0,1], None).
+    Если модели нет/ошибка: (None, 'сообщение').
+    """
     if PIPE is None:
-        return 0.0
-    proba = PIPE.predict_proba([text])[0, 1]
-    return float(proba)
-
-def infer_transformer(text: str) -> float:
-    """Вернёт P(toxic) из ruBERT-tiny2 (локальный чекпойнт)."""
+        return None, f"TF-IDF модель не загружена (ожидалась в {PIPE_PATH or 'model.joblib / artifacts/model.joblib'})."
+    try:
+        # предпочтительно predict_proba
+        if hasattr(PIPE, "predict_proba"):
+            proba = PIPE.predict_proba([text])[0, 1]
+        else:
+            # fallback: decision_function -> сигмоида
+            if hasattr(PIPE, "decision_function"):
+                z = PIPE.decision_function([text])[0]
+                proba = 1.0 / (1.0 + np.exp(-z))
+            else:
+                return None, "У модели нет predict_proba/decision_function."
+        # страховка от числ. артефактов
+        proba = float(np.clip(proba, 0.0, 1.0))
+        return proba, None
+    except Exception as e:
+        return None, f"Ошибка инференса TF-IDF: {e}"
+
+def infer_transformer(text: str):
+    """
+    Вернёт (proba, err_msg) аналогично TF-IDF.
+    """
     if TRANSFORMER["model"] is None:
-        return 0.0
+        return None, "Модель ruBERT не загружена."
     import torch
-
-    text = clean_and_stem(text)
-    if not text:
-        return 0.0
-    tok = TRANSFORMER["tokenizer"](text, return_tensors="pt", truncation=True, max_length=256)
+    text_prep = clean_and_stem(text)
+    if not text_prep:
+        return 0.0, None
+    tok = TRANSFORMER["tokenizer"](text_prep, return_tensors="pt", truncation=True, max_length=256)
     tok = {k: v.to(TRANSFORMER["device"]) for k, v in tok.items()}
-    with torch.inference_mode():
-        logits = TRANSFORMER["model"](**tok).logits
-        p = torch.softmax(logits, dim=1)[0, 1].detach().cpu().item()
-    return float(p)
+    try:
+        with torch.inference_mode():
+            logits = TRANSFORMER["model"](**tok).logits
+            proba = torch.softmax(logits, dim=1)[0, 1].detach().cpu().item()
+        return float(proba), None
+    except Exception as e:
+        return None, f"Ошибка инференса ruBERT: {e}"
 
 def predict(model_name: str, comment: str, threshold: float):
     comment = (comment or "").strip()
@@ -99,10 +147,16 @@ def predict(model_name: str, comment: str, threshold: float):
         return {"Токсичный": 0.0, "Не токсичный": 1.0}, "—"
 
     if model_name == "ruBERT-tiny2 (трансформер)":
-        p_toxic = infer_transformer(comment)
-    else:  # TF-IDF + Логистическая регрессия
-        p_toxic = infer_tfidf(comment)
+        p_toxic, err = infer_transformer(comment)
+    else:
+        p_toxic, err = infer_tfidf(comment)
+
+    if err is not None or p_toxic is None:
+        dist = {"Токсичный": 0.0, "Не токсичный": 1.0}
+        expl = f"Модель: **{model_name}**\n\n⚠️ {err}"
+        return dist, expl
 
+    # ВЕРДИКТ ТОЛЬКО ПО ЗАДАННОМУ ПОРОГУ:
     pred = "Токсичный" if p_toxic >= threshold else "Не токсичный"
     dist = {"Токсичный": p_toxic, "Не токсичный": 1 - p_toxic}
     expl = (
@@ -116,7 +170,9 @@ def predict(model_name: str, comment: str, threshold: float):
 def clear_all():
     return "ruBERT-tiny2 (трансформер)", "", DEFAULT_THRESHOLD, {"Токсичный": 0.0, "Не токсичный": 1.0}, "—"
 
-# ====== UI ======
+# ==============================
+# UI
+# ==============================
 TITLE = "Анализатор токсичности (две модели)"
 DESCRIPTION = "Выберите модель, задайте порог (по умолчанию 0.70) и введите комментарий."
 
@@ -127,20 +183,19 @@ CUSTOM_CSS = """
 
 ABOUT_MD = """
 ### Параметры и описание моделей
-
 **1) ruBERT-tiny2 (трансформер)**
 - База: `cointegrated/rubert-tiny2` (BERT-tiny для русского).
 - Токенизация: BERT WordPiece.
 - Предобработка: удаление пунктуации, нормализация спец-токенов (`url`, `tag`, `num`), стемминг Snowball.
 - Обучение: 10 эпох с early stopping (по macro-F1), class weights (balanced).
-- Рекомендованный порог по валидации: ~**0.70**.
+- Рекомендованный порог: ~**0.70**.
 
 **2) TF-IDF + Логистическая регрессия**
-- Векторизация: `TfidfVectorizer(analyzer="char_wb", ngram_range=(4,5), max_features=200k, min_df≈1.75e-4, max_df≈0.96)`.
-- Классификатор: `LogisticRegression(penalty="l1", solver="liblinear", C≈5.52, class_weight="balanced", max_iter=5000, tol≈2.4e-4)`.
-- Рекомендованный порог (по ранее полученным метрикам): ~**0.40**.
+- `TfidfVectorizer(analyzer="char_wb", ngram_range=(4,5), max_features≈200k, min_df≈1.75e-4, max_df≈0.96)`.
+- `LogisticRegression(penalty="l1", solver="liblinear", C≈5.52, class_weight="balanced", max_iter=5000, tol≈2.4e-4)`.
+- Рекомендованный порог: ~**0.40**.
 
-**Порог** можно свободно менять слайдером — выберите баланс precision/recall под задачу.
+> Порог можно свободно менять слайдером — подбирайте нужный баланс precision/recall.
 """
 
 with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="sky"), css=CUSTOM_CSS) as demo:
@@ -155,7 +210,8 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="sky"), cs
                 label="Модель"
             )
             comment_input = gr.Textbox(label="Текст комментария", lines=6, placeholder="Напишите что-нибудь…")
-            thr = gr.Slider(label="Порог классификации", minimum=0.0, maximum=1.0, value=DEFAULT_THRESHOLD, step=0.01)
+            thr = gr.Slider(label="Порог классификации", minimum=0.0, maximum=1.0,
+                            value=DEFAULT_THRESHOLD, step=0.01)
             with gr.Row():
                 analyze_btn = gr.Button("Анализ", variant="primary")
                 clear_btn = gr.Button("Очистить", variant="secondary")