app.py

import json
import re
import numpy as np
import torch
import gradio as gr
import httpx
from xml.etree import ElementTree as ET
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# ── Константы ──────────────────────────────────────────────────────────────
MODEL_DIR = "model_best"
LABEL_MAP = "data/label_mapping.json"
MAX_LEN   = 256
THRESHOLD = 0.95

CATEGORY_NAMES = {
    "cs":      "Computer Science",
    "math":    "Mathematics",
    "physics": "Physics",
    "q-bio":   "Quantitative Biology",
    "stat":    "Statistics",
    "econ":    "Economics",
    "eess":    "Electrical Engineering",
    "q-fin":   "Quantitative Finance",
}

# ── Загрузка модели ────────────────────────────────────────────────────────
def load_model():
    with open(LABEL_MAP) as f:
        label_info = json.load(f)
    id2cat = {int(k): v for k, v in label_info["id2cat"].items()}
    device = torch.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()
    return tokenizer, model, id2cat, device

tokenizer, model, id2cat, device = load_model()

# ── Утилиты ────────────────────────────────────────────────────────────────
def clean_category(cat: str) -> str:
    match = re.search(r"term':\s*'(\w[\w-]*)", cat)
    if match:
        return match.group(1)
    return cat.strip()


def extract_arxiv_id(url: str) -> str | None:
    """Вытаскивает arxiv id из ссылки любого формата."""
    # Поддерживаем форматы:
    # https://arxiv.org/abs/2301.00001
    # https://arxiv.org/pdf/2301.00001
    # 2301.00001
    match = re.search(r"(\d{4}\.\d{4,5}(v\d+)?)", url)
    return match.group(1) if match else None


def fetch_arxiv(url: str):
    arxiv_id = extract_arxiv_id(url.strip())
    if not arxiv_id:
        return None, None, "❌ Не удалось распознать arxiv ID. Пример: https://arxiv.org/abs/1706.03762"

    try:
        api_url = f"https://export.arxiv.org/api/query?id_list={arxiv_id}"
        response = httpx.get(api_url, timeout=8.0)  # жёсткий таймаут 8 сек
        response.raise_for_status()

        ns = {"atom": "http://www.w3.org/2005/Atom"}
        root = ET.fromstring(response.text)
        entry = root.find("atom:entry", ns)

        if entry is None:
            return None, None, f"❌ Статья {arxiv_id} не найдена."

        title    = entry.find("atom:title", ns).text.strip().replace("\n", " ")
        abstract = entry.find("atom:summary", ns).text.strip().replace("\n", " ")
        return title, abstract, None

    except httpx.TimeoutException:
        return None, None, "❌ arxiv не отвечает — попробуй ещё раз."
    except Exception as e:
        return None, None, f"❌ Ошибка: {e}"

# ── Предсказание ───────────────────────────────────────────────────────────
def predict(title: str, abstract: str):
    if not title or not title.strip():
        return "⚠️ Введите название статьи.", None

    enc = tokenizer(
        title.strip(),
        abstract.strip() if abstract else "",
        max_length=MAX_LEN,
        truncation=True,
        padding="max_length",
        return_tensors="pt",
    )
    with torch.no_grad():
        logits = model(
            input_ids=enc["input_ids"].to(device),
            attention_mask=enc["attention_mask"].to(device),
        ).logits

    probs = torch.softmax(logits, dim=-1).squeeze().cpu().numpy()
    sorted_ids = np.argsort(probs)[::-1]

    results, cumsum = [], 0.0
    for idx in sorted_ids:
        cat = clean_category(id2cat[idx])
        results.append((cat, float(probs[idx])))
        cumsum += probs[idx]
        if cumsum >= THRESHOLD:
            break

    lines = []
    if not abstract or not abstract.strip():
        lines.append("ℹ️ Abstract не введён — классификация только по названию.\n")

    for i, (cat, prob) in enumerate(results):
        name = CATEGORY_NAMES.get(cat, cat)
        prefix = "🥇" if i == 0 else ("🥈" if i == 1 else "▪️")
        lines.append(f"{prefix} **{name}** (`{cat}`) — {prob:.1%}")

    total = sum(p for _, p in results)
    lines.append(f"\n_Суммарная вероятность: {total:.1%}_")

    all_sorted = [(clean_category(id2cat[i]), float(probs[i])) for i in np.argsort(probs)[::-1]]
    plot_data  = {
        "Category":    [CATEGORY_NAMES.get(c, c) for c, _ in all_sorted],
        "Probability": [round(p, 4) for _, p in all_sorted],
    }

    return "\n".join(lines), plot_data


def predict_from_url(url: str):
    """Парсит arxiv и классифицирует."""
    if not url or not url.strip():
        return "", "", "⚠️ Введите ссылку на статью.", None

    title, abstract, error = fetch_arxiv(url)
    if error:
        return "", "", error, None

    result_text, plot_data = predict(title, abstract)
    return title, abstract, result_text, plot_data


# ── Интерфейс ──────────────────────────────────────────────────────────────
with gr.Blocks(title="arXiv Classifier") as demo:
    gr.Markdown(
        """
        # 📄 arXiv Article Classifier
        Определяет тематику научной статьи по названию и аннотации.
        Показывает категории с суммарной вероятностью **≥ 95%**.
        """
    )

    with gr.Tabs():

        # ── Вкладка 1: ввод вручную ────────────────────────────────────────
        with gr.Tab("✏️ Ввод вручную"):
            with gr.Row():
                with gr.Column(scale=2):
                    title_input = gr.Textbox(
                        label="Название статьи",
                        placeholder="Например: Attention Is All You Need",
                    )
                    abstract_input = gr.Textbox(
                        label="Abstract (необязательно)",
                        placeholder="Вставьте аннотацию статьи сюда…",
                        lines=7,
                    )
                    btn_manual = gr.Button("Классифицировать", variant="primary")

                with gr.Column(scale=3):
                    result_text_manual = gr.Markdown(label="Результат")
                    result_plot_manual = gr.BarPlot(
                        x="Category",
                        y="Probability",
                        title="Вероятности по категориям",
                        y_lim=[0, 1],
                        tooltip=["Category", "Probability"],
                    )

            btn_manual.click(
                fn=predict,
                inputs=[title_input, abstract_input],
                outputs=[result_text_manual, result_plot_manual],
            )

            gr.Examples(
                examples=[
                    ["Attention Is All You Need",
                     "We propose a new simple network architecture, the Transformer, based solely on attention mechanisms."],
                    ["A proof of the Riemann Hypothesis using spectral methods", ""],
                    ["CRISPR-Cas9 genome editing in mammalian cells",
                     "We demonstrate efficient genome editing in human cells using the CRISPR-Cas9 system."],
                ],
                inputs=[title_input, abstract_input],
                label="Примеры",
            )

        # ── Вкладка 2: ссылка на arxiv ─────────────────────────────────────
        with gr.Tab("🔗 Ссылка на arxiv"):
            with gr.Row():
                with gr.Column(scale=2):
                    url_input = gr.Textbox(
                        label="Ссылка на статью",
                        placeholder="https://arxiv.org/abs/1706.03762",
                    )
                    btn_url = gr.Button("Загрузить и классифицировать", variant="primary")
                    parsed_title    = gr.Textbox(label="Название (спарсено)", interactive=False)
                    parsed_abstract = gr.Textbox(label="Abstract (спарсен)", lines=5, interactive=False)

                with gr.Column(scale=3):
                    result_text_url = gr.Markdown(label="Результат")
                    result_plot_url = gr.BarPlot(
                        x="Category",
                        y="Probability",
                        title="Вероятности по категориям",
                        y_lim=[0, 1],
                        tooltip=["Category", "Probability"],
                    )

            btn_url.click(
                fn=predict_from_url,
                inputs=[url_input],
                outputs=[parsed_title, parsed_abstract, result_text_url, result_plot_url],
            )

            gr.Examples(
                examples=[
                    ["https://arxiv.org/abs/1706.03762"],  # Attention is all you need
                    ["https://arxiv.org/abs/1810.04805"],  # BERT
                    ["https://arxiv.org/abs/2303.08774"],  # GPT-4
                ],
                inputs=[url_input],
                label="Примеры ссылок",
            )

    with gr.Accordion("Поддерживаемые категории", open=False):
        gr.Markdown("\n".join(f"- `{k}` — {v}" for k, v in CATEGORY_NAMES.items()))


if __name__ == "__main__":
    demo.launch()