Upload article_classifier_training.ipynb

notebooks/article_classifier_training.ipynb

@@ -0,0 +1,231 @@
+{
+ "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
+}