{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "056645ed",
   "metadata": {},
   "source": [
    "# Обучение классификатора тематик статей\n",
    "\n",
    "Этот ноутбук повторяет логику `src/train.py`, но в формате, удобном для сдачи и экспериментов.\n",
    "Перед запуском скачайте датасет и положите его в `data/`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7563cfb3",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q transformers datasets accelerate evaluate scikit-learn pandas pyarrow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fc396eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "DATA_PATH = Path(\"data/arxiv_data.csv\")\n",
    "MODEL_NAME = \"allenai/scibert_scivocab_uncased\"\n",
    "OUTPUT_DIR = Path(\"artifacts/article_topic_model\")\n",
    "MAX_LENGTH = 256\n",
    "TEST_SIZE = 0.15\n",
    "SEED = 42"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a84a37ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "from src.data_utils import load_dataset_frame, filter_rare_classes\n",
    "\n",
    "df = load_dataset_frame(DATA_PATH, text_cols=[\"title\", \"abstract\"], label_col=\"tag\")\n",
    "df = filter_rare_classes(df, min_examples_per_class=20)\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b02568e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "df[\"label\"].value_counts().head(20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9721b73f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "labels = sorted(df[\"label\"].unique().tolist())\n",
    "label2id = {label: idx for idx, label in enumerate(labels)}\n",
    "id2label = {idx: label for label, idx in label2id.items()}\n",
    "df[\"label_id\"] = df[\"label\"].map(label2id)\n",
    "\n",
    "train_df, valid_df = train_test_split(\n",
    "    df,\n",
    "    test_size=TEST_SIZE,\n",
    "    random_state=SEED,\n",
    "    stratify=df[\"label_id\"],\n",
    ")\n",
    "\n",
    "len(train_df), len(valid_df), len(labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11dc01f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import evaluate\n",
    "from datasets import Dataset, DatasetDict\n",
    "from transformers import (\n",
    "    AutoTokenizer,\n",
    "    AutoModelForSequenceClassification,\n",
    "    DataCollatorWithPadding,\n",
    "    Trainer,\n",
    "    TrainingArguments,\n",
    ")\n",
    "\n",
    "ds = DatasetDict({\n",
    "    \"train\": Dataset.from_pandas(train_df[[\"text\", \"label_id\"]], preserve_index=False),\n",
    "    \"validation\": Dataset.from_pandas(valid_df[[\"text\", \"label_id\"]], preserve_index=False),\n",
    "})\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)\n",
    "\n",
    "def tokenize_batch(batch):\n",
    "    return tokenizer(batch[\"text\"], truncation=True, max_length=MAX_LENGTH)\n",
    "\n",
    "tokenized = ds.map(tokenize_batch, batched=True)\n",
    "tokenized = tokenized.rename_column(\"label_id\", \"labels\")\n",
    "tokenized.set_format(type=\"torch\", columns=[\"input_ids\", \"attention_mask\", \"labels\"])\n",
    "\n",
    "model = AutoModelForSequenceClassification.from_pretrained(\n",
    "    MODEL_NAME,\n",
    "    num_labels=len(labels),\n",
    "    id2label={int(k): v for k, v in id2label.items()},\n",
    "    label2id=label2id,\n",
    ")\n",
    "\n",
    "accuracy_metric = evaluate.load(\"accuracy\")\n",
    "f1_metric = evaluate.load(\"f1\")\n",
    "\n",
    "def compute_metrics(eval_pred):\n",
    "    logits, labels_np = eval_pred\n",
    "    preds = np.argmax(logits, axis=-1)\n",
    "    return {\n",
    "        \"accuracy\": accuracy_metric.compute(predictions=preds, references=labels_np)[\"accuracy\"],\n",
    "        \"f1_macro\": f1_metric.compute(predictions=preds, references=labels_np, average=\"macro\")[\"f1\"],\n",
    "        \"f1_weighted\": f1_metric.compute(predictions=preds, references=labels_np, average=\"weighted\")[\"f1\"],\n",
    "    }\n",
    "\n",
    "args = TrainingArguments(\n",
    "    output_dir=str(OUTPUT_DIR / \"checkpoints\"),\n",
    "    learning_rate=2e-5,\n",
    "    per_device_train_batch_size=8,\n",
    "    per_device_eval_batch_size=16,\n",
    "    num_train_epochs=3,\n",
    "    weight_decay=0.01,\n",
    "    evaluation_strategy=\"epoch\",\n",
    "    save_strategy=\"epoch\",\n",
    "    logging_steps=50,\n",
    "    load_best_model_at_end=True,\n",
    "    metric_for_best_model=\"f1_macro\",\n",
    "    report_to=\"none\",\n",
    ")\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=args,\n",
    "    train_dataset=tokenized[\"train\"],\n",
    "    eval_dataset=tokenized[\"validation\"],\n",
    "    tokenizer=tokenizer,\n",
    "    data_collator=DataCollatorWithPadding(tokenizer=tokenizer),\n",
    "    compute_metrics=compute_metrics,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0f53c46",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17552445",
   "metadata": {},
   "outputs": [],
   "source": [
    "metrics = trainer.evaluate()\n",
    "metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "702cda86",
   "metadata": {},
   "outputs": [],
   "source": [
    "OUTPUT_DIR.mkdir(parents=True, exist_ok=True)\n",
    "trainer.model.save_pretrained(OUTPUT_DIR)\n",
    "tokenizer.save_pretrained(OUTPUT_DIR)\n",
    "\n",
    "import json\n",
    "(OUTPUT_DIR / \"label_mapping.json\").write_text(\n",
    "    json.dumps({\n",
    "        \"label2id\": label2id,\n",
    "        \"id2label\": {str(k): v for k, v in id2label.items()},\n",
    "    }, ensure_ascii=False, indent=2),\n",
    "    encoding=\"utf-8\",\n",
    ")\n",
    "(OUTPUT_DIR / \"metrics.json\").write_text(\n",
    "    json.dumps(metrics, ensure_ascii=False, indent=2),\n",
    "    encoding=\"utf-8\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cc9ca80",
   "metadata": {},
   "outputs": [],
   "source": [
    "from src.predictor import ArticleTopicPredictor\n",
    "\n",
    "predictor = ArticleTopicPredictor(OUTPUT_DIR)\n",
    "predictor.predict(\n",
    "    title=\"Graph neural networks for molecular property prediction\",\n",
    "    abstract=\"We benchmark graph neural architectures on several chemistry tasks.\",\n",
    ")"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 5
}