Update app.py

app.py

@@ -1,112 +1,134 @@
-import gradio as gr
-import joblib, json, re
+import os, json, re
 from html import unescape
 
+import gradio as gr
+import torch
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from nltk.stem.snowball import RussianStemmer
+
 # -----------------------------
-# 1) Точная копия preprocessor (без изменений)
+# 1) Предобработка: как в обучении
 # -----------------------------
-_URL_RE = re.compile(r'https?://\S+|www\.\S+')
-_TAG_RE = re.compile(r'[@#]\w+')
-_NUM_RE = re.compile(r'\d+')
-_WS_RE  = re.compile(r'\s+')
+_URL_RE   = re.compile(r'https?://\S+|www\.\S+')
+_TAG_RE   = re.compile(r'[@#]\w+')
+_NUM_RE   = re.compile(r'\d+')
+_PUNCT_RE = re.compile(r"[^\w\s]+", flags=re.UNICODE)
+_WS_RE    = re.compile(r"\s+")
 
-def clean_text(s: str) -> str:
-    """ДОЛЖНА совпадать с версией из обучения, иначе pickle не найдёт функцию."""
+stemmer = RussianStemmer(ignore_stopwords=False)
+
+def clean_and_stem(s: str) -> str:
     if not isinstance(s, str):
         s = str(s)
     s = unescape(s).lower()
-    s = _URL_RE.sub(' <url> ', s)
-    s = _TAG_RE.sub(' <tag> ', s)
-    s = _NUM_RE.sub(' <num> ', s)
-    s = s.replace('\n', ' ').replace('\t', ' ')
-    s = _WS_RE.sub(' ', s).strip()
-    return s
-
-# ---------------------------------
-# 2) Загрузка пайплайна (без конфига)
-# ---------------------------------
-PIPE = joblib.load("model.joblib")
-
-# ---------------------------------
-# 3) Максимально упрощенный инференс
-# ---------------------------------
-def predict(comment: str):
+    s = _URL_RE.sub(" url ", s)
+    s = _TAG_RE.sub(" tag ", s)
+    s = _NUM_RE.sub(" num ", s)
+    s = _PUNCT_RE.sub(" ", s)
+    s = _WS_RE.sub(" ", s).strip()
+    if not s:
+        return s
+    out = []
+    for t in s.split():
+        out.append(t if t in {"url", "tag", "num"} else stemmer.stem(t))
+    return " ".join(out)
+
+# -----------------------------
+# 2) Загрузка модели/токенайзера
+# -----------------------------
+# Папка с файлами модели; по умолчанию 'best', можно переопределить переменной окружения MODEL_DIR
+MODEL_DIR = os.getenv("MODEL_DIR", "best" if os.path.exists("best/config.json") else ".")
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
+model = AutoModelForSequenceClassification.from_pretrained(MODEL_DIR)
+model.to(DEVICE).eval()
+
+# Порог по умолчанию берём из inference_config.json (если есть), иначе 0.40
+DEFAULT_THRESHOLD = 0.40
+try:
+    with open(os.path.join(MODEL_DIR, "inference_config.json"), "r", encoding="utf-8") as f:
+        DEFAULT_THRESHOLD = float(json.load(f).get("threshold_val", DEFAULT_THRESHOLD))
+except Exception:
+    pass
+
+MAX_LEN = 256
+
+@torch.no_grad()
+def predict(text: str, threshold: float):
     """
-    Возвращает словарь {метка: вероятность} для компонента gr.Label.
+    Возвращает:
+      - dict для gr.Label с распределением вероятностей
+      - пояснение (Markdown)
     """
-    if not comment or not comment.strip():
-        return None # Возвращаем None, чтобы очистить поле вывода
-    
-    proba_toxic = float(PIPE.predict_proba([comment])[0, 1])
-    proba_not_toxic = 1 - proba_toxic
-    
-    # Возвращаем словарь, gr.Label сам подсветит класс с большей вероятностью
-    return {"Токсичный": proba_toxic, "Не токсичный": proba_not_toxic}
-        
-# ---------------------------------
-# 4) Минималистичный интерфейс
-# ---------------------------------
-
-TITLE = "Анализатор токсичности комментариев"
-DESCRIPTION = "Введите комментарий на русском языке. Модель покажет распределение вероятностей между классами «Токсичный» и «Не токсичный»."
+    if not text or not text.strip():
+        return None, "Введите текст выше."
+
+    text_prep = clean_and_stem(text)
+    batch = tokenizer(text_prep, truncation=True, max_length=MAX_LEN,
+                      padding=True, return_tensors="pt")
+    batch = {k: v.to(DEVICE) for k, v in batch.items()}
+
+    logits = model(**batch).logits
+    probs = torch.softmax(logits, dim=1).cpu().numpy()[0]
+    p_not, p_tox = float(probs[0]), float(probs[1])
+
+    label = "Токсичный" if p_tox >= threshold else "Не токсичный"
+    md = (
+        f"**Порог:** {threshold:.2f}  \n"
+        f"**Класс:** **{label}**  \n"
+        f"P(toxic) = {p_tox:.3f} · P(not_toxic) = {p_not:.3f}"
+    )
+    return {"Токсичный": p_tox, "Не ток��ичный": p_not}, md
+
+# -----------------------------
+# 3) UI
+# -----------------------------
+TITLE = "Анализатор токсичности комментариев (ruBERT-tiny2)"
+DESCRIPTION = (
+    "Введите комментарий на русском языке. "
+    "Модель (cointegrated/rubert-tiny2, дообученная на вашем датасете) "
+    "вернёт распределение вероятностей и итоговую метку с учётом порога."
+)
 ARTICLE = """
----
-### Детали модели
-*   **Архитектура**: TF-IDF (char_wb, n-граммы 4-5) + Логистическая регрессия (L1, class_weight=balanced).
-*   **Назначение**: Классификация русскоязычных текстов.
+### Детали
+- Архитектура: **ruBERT-tiny2** → linear head (2 класса).  
+- Предобработка: замена URL/тегов/чисел, удаление пунктуации, **стемминг** (NLTK).  
+- Ввод порога позволяет управлять балансом Precision/Recall.
 """
 
-# CSS для подключения и применения шрифта "Inter"
 CUSTOM_CSS = """
 @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;700&display=swap');
-gradio-app { 
-    font-family: 'Inter', sans-serif;
-}
+gradio-app { font-family: 'Inter', sans-serif; }
 """
 
-with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="sky"), css=CUSTOM_CSS) as demo:
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="sky"),
+               css=CUSTOM_CSS, flagging_mode="never") as demo:
     gr.Markdown(f"# {TITLE}")
     gr.Markdown(DESCRIPTION)
 
     with gr.Row():
-        # Левая колонка для ввода
         with gr.Column(scale=2):
-            comment_input = gr.Textbox(
-                label="Текст комментария",
-                lines=5,
-                placeholder="Напишите что-нибудь...",
-            )
+            txt = gr.Textbox(label="Текст комментария", lines=6,
+                             placeholder="Напишите что-нибудь…")
+            thr = gr.Slider(0.0, 1.0, value=DEFAULT_THRESHOLD, step=0.01,
+                            label="Порог классификации")
             with gr.Row():
-                clear_btn = gr.Button("Очистить", variant="secondary")
-                analyze_btn = gr.Button("Анализ", variant="primary")
+                btn_clear = gr.Button("Очистить", variant="secondary")
+                btn_pred  = gr.Button("Анализ", variant="primary")
 
-        # Правая колонка для вывода
         with gr.Column(scale=1):
-            result_label = gr.Label(label="Результат", num_top_classes=2)
-            
+            dist = gr.Label(label="Распределение вероятностей", num_top_classes=2)
+            info = gr.Markdown()
+
     gr.Markdown(ARTICLE)
-    
-    # --- Логика взаимодействия компонентов ---
-    
-    def clear_all():
-        return "", None
-
-    # Привязка функций к кнопкам и событиям
-    analyze_btn.click(
-        fn=predict,
-        inputs=comment_input,
-        outputs=result_label
-    )
-    comment_input.submit(
-        fn=predict,
-        inputs=comment_input,
-        outputs=result_label
-    )
-    clear_btn.click(
-        fn=clear_all, 
-        inputs=[], 
-        outputs=[comment_input, result_label]
-    )
+
+    def _clear():
+        return "", DEFAULT_THRESHOLD, None, " "
+
+    btn_pred.click(predict, inputs=[txt, thr], outputs=[dist, info])
+    txt.submit(predict, inputs=[txt, thr], outputs=[dist, info])
+    btn_clear.click(_clear, outputs=[txt, thr, dist, info])
 
 if __name__ == "__main__":
     demo.launch()