Synced repo using 'sync_with_huggingface' Github Action

notebooks/mamba_vs_transformerts.ipynb

@@ -0,0 +1,361 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import time\n",
+    "import torch\n",
+    "import warnings\n",
+    "import numpy as np\n",
+    "import pandas as pd\n",
+    "import torch.nn as nn\n",
+    "import torch.optim as optim\n",
+    "from torch.utils.data import DataLoader, TensorDataset\n",
+    "from sklearn.metrics import accuracy_score, precision_recall_fscore_support\n",
+    "for warn in [UserWarning, FutureWarning]: warnings.filterwarnings(\"ignore\", category = warn)\n",
+    "\n",
+    "from src.data_utils.config import DatasetConfig\n",
+    "from src.data_utils.dataset_params import DatasetName\n",
+    "from src.data_utils.dataset_generator import DatasetGenerator\n",
+    "from src.models.models import TransformerClassifier, CustomMambaClassifier, LSTMClassifier\n",
+    "\n",
+    "MAX_SEQ_LEN = 300\n",
+    "EMBEDDING_DIM = 128\n",
+    "BATCH_SIZE = 32\n",
+    "LEARNING_RATE = 7e-5 # уменьшили lr: 1e-4 -> 7e-5\n",
+    "NUM_EPOCHS = 20 # подняли количество эпох: 5 -> 20\n",
+    "NUM_CLASSES = 2\n",
+    "\n",
+    "SAVE_DIR = \"../best_models/\"\n",
+    "os.makedirs(SAVE_DIR, exist_ok=True)\n",
+    "DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+    "\n",
+    "config = DatasetConfig(\n",
+    "    load_from_disk=True,\n",
+    "    path_to_data=\"../datasets\",\n",
+    "    train_size=25000,  # увеличили количество сэмплов\n",
+    "    val_size=12500,\n",
+    "    test_size=12500\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "generator = DatasetGenerator(DatasetName.IMDB, config=config)\n",
+    "(X_train, y_train), (X_val, y_val), (X_test, y_test) = generator.generate_dataset()\n",
+    "VOCAB_SIZE = len(generator.vocab)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Также создадим генератор для тестовых данных второго датасета"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "text_processor = generator.get_text_processor()\n",
+    "config_polarity = DatasetConfig(\n",
+    "    load_from_disk=True,\n",
+    "    path_to_data=\"../datasets\",\n",
+    "    test_size=10000,\n",
+    "    build_vocab=False\n",
+    ")\n",
+    "generator_polarity = DatasetGenerator(DatasetName.POLARITY, config=config_polarity)\n",
+    "generator_polarity.vocab = generator.vocab\n",
+    "generator_polarity.id2word = generator.id2word\n",
+    "generator_polarity.text_processor = text_processor\n",
+    "(_, _), (_, _), (X_test_polarity, y_test_polarity) = generator_polarity.generate_dataset()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "def train_and_evaluate(model, train_loader, val_loader, optimizer, criterion, num_epochs, device, model_name, save_path):\n",
+    "    best_val_f1 = 0.0\n",
+    "    history = {'train_loss': [], 'val_loss': [], 'val_accuracy': [], 'val_f1': []}\n",
+    "    \n",
+    "    print(f\"--- Начало обучения модели: {model_name} на устройстве {device} ---\")\n",
+    "\n",
+    "    for epoch in range(num_epochs):\n",
+    "        model.train()\n",
+    "        start_time = time.time()\n",
+    "        total_train_loss = 0\n",
+    "\n",
+    "        for batch_X, batch_y in train_loader:\n",
+    "            batch_X, batch_y = batch_X.to(device), batch_y.to(device)\n",
+    "            optimizer.zero_grad()\n",
+    "            outputs = model(batch_X)\n",
+    "            loss = criterion(outputs, batch_y)\n",
+    "            loss.backward()\n",
+    "            optimizer.step()\n",
+    "            total_train_loss += loss.item()\n",
+    "        \n",
+    "        avg_train_loss = total_train_loss / len(train_loader)\n",
+    "        history['train_loss'].append(avg_train_loss)\n",
+    "\n",
+    "        model.eval()\n",
+    "        total_val_loss = 0\n",
+    "        all_preds = []\n",
+    "        all_labels = []\n",
+    "\n",
+    "        with torch.no_grad():\n",
+    "            for batch_X, batch_y in val_loader:\n",
+    "                batch_X, batch_y = batch_X.to(device), batch_y.to(device)\n",
+    "                outputs = model(batch_X)\n",
+    "                loss = criterion(outputs, batch_y)\n",
+    "                total_val_loss += loss.item()\n",
+    "                \n",
+    "                _, predicted = torch.max(outputs.data, 1)\n",
+    "                all_preds.extend(predicted.cpu().numpy())\n",
+    "                all_labels.extend(batch_y.cpu().numpy())\n",
+    "        \n",
+    "        avg_val_loss = total_val_loss / len(val_loader)\n",
+    "        \n",
+    "        accuracy = accuracy_score(all_labels, all_preds)\n",
+    "        precision, recall, f1, _ = precision_recall_fscore_support(all_labels, all_preds, average='binary')\n",
+    "        \n",
+    "        history['val_loss'].append(avg_val_loss)\n",
+    "        history['val_accuracy'].append(accuracy)\n",
+    "        history['val_f1'].append(f1)\n",
+    "\n",
+    "        epoch_time = time.time() - start_time\n",
+    "        print(f\"Эпоха {epoch+1}/{num_epochs} | Время: {epoch_time:.2f}с | Train Loss: {avg_train_loss:.4f} | \"\n",
+    "              f\"Val Loss: {avg_val_loss:.4f} | Val Acc: {accuracy:.4f} | Val F1: {f1:.4f}\")\n",
+    "\n",
+    "        if f1 > best_val_f1:\n",
+    "            best_val_f1 = f1\n",
+    "            torch.save(model.state_dict(), save_path)\n",
+    "            print(f\"  -> Модель сохранена, новый лучший Val F1: {best_val_f1:.4f}\")\n",
+    "            \n",
+    "    print(f\"--- Обучение модели {model_name} завершено ---\")\n",
+    "    return history\n",
+    "\n",
+    "def evaluate_on_test(model, test_loader, device, criterion):\n",
+    "    model.eval()\n",
+    "    total_test_loss = 0\n",
+    "    all_preds = []\n",
+    "    all_labels = []\n",
+    "\n",
+    "    with torch.no_grad():\n",
+    "        for batch_X, batch_y in test_loader:\n",
+    "            batch_X, batch_y = batch_X.to(device), batch_y.to(device)\n",
+    "            outputs = model(batch_X)\n",
+    "            loss = criterion(outputs, batch_y)\n",
+    "            total_test_loss += loss.item()\n",
+    "            \n",
+    "            _, predicted = torch.max(outputs.data, 1)\n",
+    "            all_preds.extend(predicted.cpu().numpy())\n",
+    "            all_labels.extend(batch_y.cpu().numpy())\n",
+    "            \n",
+    "    avg_test_loss = total_test_loss / len(test_loader)\n",
+    "        \n",
+    "    accuracy = accuracy_score(all_labels, all_preds)\n",
+    "    precision, recall, f1, _ = precision_recall_fscore_support(all_labels, all_preds, average='binary')\n",
+    "    \n",
+    "    return {'loss': avg_test_loss, 'accuracy': accuracy, 'precision': precision, 'recall': recall, 'f1_score': f1}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "def create_dataloader(X, y, batch_size, shuffle=True):\n",
+    "    X_tensor = torch.as_tensor(X, dtype=torch.long)\n",
+    "    y_tensor = torch.as_tensor(y, dtype=torch.long)\n",
+    "    dataset = TensorDataset(X_tensor, y_tensor)\n",
+    "    return DataLoader(dataset, batch_size=batch_size, shuffle=shuffle)\n",
+    "\n",
+    "train_loader = create_dataloader(X_train, y_train, BATCH_SIZE)\n",
+    "val_loader = create_dataloader(X_val, y_val, BATCH_SIZE, shuffle=False)\n",
+    "test_loader = create_dataloader(X_test, y_test, BATCH_SIZE, shuffle=False)\n",
+    "test_loader_polarity = create_dataloader(X_test_polarity, y_test_polarity, BATCH_SIZE, shuffle=False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "--- Начало обучения модели: CustomMamba на устройстве cuda ---\n",
+      "Эпоха 1/20 | Время: 897.35с | Train Loss: 0.6261 | Val Loss: 0.5313 | Val Acc: 0.7388 | Val F1: 0.7503\n",
+      "  -> Модель сохранена, новый лучший Val F1: 0.7503\n",
+      "Эпоха 2/20 | Время: 903.31с | Train Loss: 0.4748 | Val Loss: 0.4559 | Val Acc: 0.7889 | Val F1: 0.7901\n",
+      "  -> Модель сохранена, новый лучший Val F1: 0.7901\n",
+      "Эпоха 3/20 | Время: 930.96с | Train Loss: 0.3955 | Val Loss: 0.4176 | Val Acc: 0.8090 | Val F1: 0.8048\n",
+      "  -> Модель сохранена, новый лучший Val F1: 0.8048\n",
+      "Эпоха 4/20 | Время: 952.79с | Train Loss: 0.3429 | Val Loss: 0.3998 | Val Acc: 0.8230 | Val F1: 0.8303\n",
+      "  -> Модель сохранена, новый лучший Val F1: 0.8303\n",
+      "Эпоха 5/20 | Время: 904.67с | Train Loss: 0.2984 | Val Loss: 0.4387 | Val Acc: 0.8165 | Val F1: 0.8337\n",
+      "  -> Модель сохранена, новый лучший Val F1: 0.8337\n",
+      "Эпоха 6/20 | Время: 935.58с | Train Loss: 0.2609 | Val Loss: 0.4219 | Val Acc: 0.8255 | Val F1: 0.8386\n",
+      "  -> Модель сохранена, новый лучший Val F1: 0.8386\n",
+      "Эпоха 7/20 | Время: 902.73с | Train Loss: 0.2266 | Val Loss: 0.4342 | Val Acc: 0.8334 | Val F1: 0.8420\n",
+      "  -> Модель сохранена, новый лучший Val F1: 0.8420\n",
+      "Эпоха 8/20 | Время: 921.49с | Train Loss: 0.1953 | Val Loss: 0.4675 | Val Acc: 0.8299 | Val F1: 0.8411\n",
+      "Эпоха 9/20 | Время: 892.54с | Train Loss: 0.1590 | Val Loss: 0.5205 | Val Acc: 0.8359 | Val F1: 0.8394\n",
+      "Эпоха 10/20 | Время: 893.22с | Train Loss: 0.1263 | Val Loss: 0.6014 | Val Acc: 0.8303 | Val F1: 0.8398\n",
+      "Эпоха 11/20 | Время: 953.69с | Train Loss: 0.0945 | Val Loss: 0.8025 | Val Acc: 0.8183 | Val F1: 0.8352\n",
+      "Эпоха 12/20 | Время: 924.98с | Train Loss: 0.0644 | Val Loss: 0.8539 | Val Acc: 0.8290 | Val F1: 0.8369\n",
+      "Эпоха 13/20 | Время: 916.92с | Train Loss: 0.0428 | Val Loss: 1.0646 | Val Acc: 0.8286 | Val F1: 0.8266\n",
+      "Эпоха 14/20 | Время: 904.39с | Train Loss: 0.0266 | Val Loss: 1.5225 | Val Acc: 0.8149 | Val F1: 0.8305\n",
+      "Эпоха 15/20 | Время: 923.80с | Train Loss: 0.0199 | Val Loss: 1.6176 | Val Acc: 0.8242 | Val F1: 0.8351\n",
+      "Эпоха 16/20 | Время: 922.03с | Train Loss: 0.0134 | Val Loss: 1.8983 | Val Acc: 0.8258 | Val F1: 0.8354\n",
+      "Эпоха 17/20 | Время: 914.58с | Train Loss: 0.0069 | Val Loss: 1.7992 | Val Acc: 0.8260 | Val F1: 0.8252\n",
+      "Эпоха 18/20 | Время: 937.09с | Train Loss: 0.0134 | Val Loss: 2.2935 | Val Acc: 0.8120 | Val F1: 0.8287\n",
+      "Эпоха 19/20 | Время: 898.26с | Train Loss: 0.0061 | Val Loss: 2.3201 | Val Acc: 0.8294 | Val F1: 0.8338\n",
+      "Эпоха 20/20 | Время: 910.46с | Train Loss: 0.0065 | Val Loss: 1.7146 | Val Acc: 0.8282 | Val F1: 0.8311\n",
+      "--- Обучение модели CustomMamba завершено ---\n",
+      "--- Оценка лучшей модели CustomMamba на тестовых данных ---\n",
+      "Результаты для CustomMamba: {'loss': 0.4370202890042301, 'accuracy': 0.8328, 'precision': 0.8017116333043226, 'recall': 0.88432, 'f1_score': 0.8409920876445527}\n",
+      "------------------------------------------------------------\n",
+      "--- Начало обучения модели: Lib_Transformer на устройстве cuda ---\n",
+      "Эпоха 1/20 | Время: 21.90с | Train Loss: 0.5894 | Val Loss: 0.5326 | Val Acc: 0.7487 | Val F1: 0.7689\n",
+      "  -> Модель сохранена, новый лучший Val F1: 0.7689\n",
+      "Эпоха 2/20 | Время: 21.56с | Train Loss: 0.4505 | Val Loss: 0.4653 | Val Acc: 0.7953 | Val F1: 0.7894\n",
+      "  -> Модель сохранена, новый лучший Val F1: 0.7894\n",
+      "Эпоха 3/20 | Время: 21.64с | Train Loss: 0.3925 | Val Loss: 0.4553 | Val Acc: 0.8001 | Val F1: 0.7818\n",
+      "Эпоха 4/20 | Время: 22.27с | Train Loss: 0.3562 | Val Loss: 0.4642 | Val Acc: 0.7836 | Val F1: 0.7480\n",
+      "Эпоха 5/20 | Время: 21.73с | Train Loss: 0.3294 | Val Loss: 0.4035 | Val Acc: 0.8305 | Val F1: 0.8337\n",
+      "  -> Модель сохранена, новый лучший Val F1: 0.8337\n",
+      "Эпоха 6/20 | Время: 21.58с | Train Loss: 0.3022 | Val Loss: 0.3936 | Val Acc: 0.8364 | Val F1: 0.8356\n",
+      "  -> Модель сохранена, новый лучший Val F1: 0.8356\n",
+      "Эпоха 7/20 | Время: 21.67с | Train Loss: 0.2752 | Val Loss: 0.3853 | Val Acc: 0.8380 | Val F1: 0.8345\n",
+      "Эпоха 8/20 | Время: 21.82с | Train Loss: 0.2507 | Val Loss: 0.3882 | Val Acc: 0.8377 | Val F1: 0.8329\n",
+      "Эпоха 9/20 | Время: 21.81с | Train Loss: 0.2286 | Val Loss: 0.4488 | Val Acc: 0.8118 | Val F1: 0.8333\n",
+      "Эпоха 10/20 | Время: 21.65с | Train Loss: 0.2056 | Val Loss: 0.3876 | Val Acc: 0.8402 | Val F1: 0.8350\n",
+      "Эпоха 11/20 | Время: 21.60с | Train Loss: 0.1803 | Val Loss: 0.3949 | Val Acc: 0.8358 | Val F1: 0.8385\n",
+      "  -> Модель сохранена, новый лучший Val F1: 0.8385\n",
+      "Эпоха 12/20 | Время: 21.57с | Train Loss: 0.1605 | Val Loss: 0.4024 | Val Acc: 0.8360 | Val F1: 0.8414\n",
+      "  -> Модель сохранена, новый лучший Val F1: 0.8414\n",
+      "Эпоха 13/20 | Время: 21.65с | Train Loss: 0.1392 | Val Loss: 0.4087 | Val Acc: 0.8356 | Val F1: 0.8340\n",
+      "Эпоха 14/20 | Время: 21.57с | Train Loss: 0.1172 | Val Loss: 0.4315 | Val Acc: 0.8323 | Val F1: 0.8297\n",
+      "Эпоха 15/20 | Время: 21.56с | Train Loss: 0.1005 | Val Loss: 0.4626 | Val Acc: 0.8317 | Val F1: 0.8284\n",
+      "Эпоха 16/20 | Время: 21.65с | Train Loss: 0.0876 | Val Loss: 0.4680 | Val Acc: 0.8318 | Val F1: 0.8335\n",
+      "Эпоха 17/20 | Время: 21.73с | Train Loss: 0.0728 | Val Loss: 0.4823 | Val Acc: 0.8317 | Val F1: 0.8326\n",
+      "Эпоха 18/20 | Время: 21.55с | Train Loss: 0.0656 | Val Loss: 0.5540 | Val Acc: 0.8206 | Val F1: 0.8068\n",
+      "Эпоха 19/20 | Время: 21.56с | Train Loss: 0.0491 | Val Loss: 0.6002 | Val Acc: 0.8235 | Val F1: 0.8178\n",
+      "Эпоха 20/20 | Время: 21.60с | Train Loss: 0.0445 | Val Loss: 0.5776 | Val Acc: 0.8314 | Val F1: 0.8318\n",
+      "--- Обучение модели Lib_Transformer завершено ---\n",
+      "--- Оценка лучшей модели Lib_Transformer на тестовых данных ---\n",
+      "Результаты для Lib_Transformer: {'loss': 0.3889380347202806, 'accuracy': 0.84488, 'precision': 0.8214765100671141, 'recall': 0.88128, 'f1_score': 0.8503280586646083}\n",
+      "------------------------------------------------------------\n",
+      "\n",
+      "\n",
+      "--- Итоговая таблица сравнения моделей на тестовых данных ---\n",
+      "                              loss  accuracy  precision    recall  f1_score\n",
+      "CustomMamba               0.437020   0.83280   0.801712  0.884320  0.840992\n",
+      "CustomMamba_polarity      0.567920   0.73850   0.688808  0.869522  0.768686\n",
+      "Lib_Transformer           0.388938   0.84488   0.821477  0.881280  0.850328\n",
+      "Lib_Transformer_polarity  0.543388   0.73980   0.690897  0.867320  0.769122\n"
+     ]
+    }
+   ],
+   "source": [
+    "model_configs = {\n",
+    "    \"CustomMamba\": {\n",
+    "        \"class\": CustomMambaClassifier,\n",
+    "        \"params\": {'vocab_size': VOCAB_SIZE, 'd_model': EMBEDDING_DIM, 'd_state': 8, \n",
+    "                   'd_conv': 4, 'num_layers': 2, 'num_classes': NUM_CLASSES},\n",
+    "    },\n",
+    "\n",
+    "    \"Lib_Transformer\": {\n",
+    "        \"class\": TransformerClassifier,\n",
+    "        \"params\": {'vocab_size': VOCAB_SIZE, 'embed_dim': EMBEDDING_DIM, 'num_heads': 8, \n",
+    "                   'num_layers': 4, 'num_classes': NUM_CLASSES, 'max_seq_len': MAX_SEQ_LEN},\n",
+    "        # num_layers: 2 -> 4\n",
+    "        # num_heads: 4 -> 8\n",
+    "    },\n",
+    "}\n",
+    "\n",
+    "results = {}\n",
+    "for model_name, config in model_configs.items():\n",
+    "\n",
+    "    model_path = os.path.join(SAVE_DIR, f\"best_model_{model_name.lower()}.pth\")\n",
+    "    \n",
+    "    model = config['class'](**config['params']).to(DEVICE)\n",
+    "    optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)\n",
+    "    criterion = nn.CrossEntropyLoss()\n",
+    "    \n",
+    "    train_and_evaluate(\n",
+    "        model=model, train_loader=train_loader, val_loader=val_loader,\n",
+    "        optimizer=optimizer, criterion=criterion, num_epochs=NUM_EPOCHS,\n",
+    "        device=DEVICE, model_name=model_name, save_path=model_path\n",
+    "    )\n",
+    "    \n",
+    "    print(f\"--- Оценка лучшей модели {model_name} на тестовых данных ---\")\n",
+    "    if os.path.exists(model_path):\n",
+    "        best_model = config['class'](**config['params']).to(DEVICE)\n",
+    "        best_model.load_state_dict(torch.load(model_path))\n",
+    "        test_metrics = evaluate_on_test(best_model, test_loader, DEVICE, criterion)\n",
+    "        results[model_name] = test_metrics\n",
+    "        results[model_name + \"_polarity\"] = evaluate_on_test(best_model, test_loader_polarity, DEVICE, criterion)\n",
+    "        print(f\"Результаты для {model_name}: {test_metrics}\")\n",
+    "    else:\n",
+    "        print(f\"Файл лучшей модели для {model_name} не найден. Пропускаем оценку.\")\n",
+    "\n",
+    "    print(\"-\" * 60)\n",
+    "    \n",
+    "if results:\n",
+    "    results_df = pd.DataFrame(results).T\n",
+    "    print(\"\\n\\n--- Итоговая таблица сравнения моделей на тестовых данных ---\")\n",
+    "    print(results_df.to_string())\n",
+    "else:\n",
+    "    print(\"Не удалось получить результаты ни для одной модели.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Видим, что transformer сильно выигрывает у мамбы как и по времени, так и по качеству. Дальше посмотрим как они справляются на данных из других датасетов"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "monkey-coding-dl-project-F4QJzkF_-py3.12",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

notebooks/models_comparations_second_dataset.ipynb

@@ -0,0 +1,240 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Восстановим константы, словарь и модели из прошлого нотубка"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import pandas as pd\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.nn.functional as F\n",
+    "from torch.utils.data import DataLoader, TensorDataset\n",
+    "\n",
+    "from src.models.models import TransformerClassifier, LSTMClassifier, CustomMambaClassifier, SimpleMambaBlock\n",
+    "from src.data_utils.config import DatasetConfig\n",
+    "from src.data_utils.dataset_params import DatasetName\n",
+    "from src.data_utils.dataset_generator import DatasetGenerator\n",
+    "\n",
+    "MAX_SEQ_LEN = 300\n",
+    "EMBEDDING_DIM = 128\n",
+    "BATCH_SIZE = 32 \n",
+    "NUM_CLASSES = 2\n",
+    "SAVE_DIR = \"../pretrained_comparison\" \n",
+    "DATA_DIR = \"../datasets\" \n",
+    "DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+    "\n",
+    "config = DatasetConfig(load_from_disk=True, path_to_data=DATA_DIR)\n",
+    "generator = DatasetGenerator(DatasetName.IMDB, config=config)\n",
+    "\n",
+    "_, _, _ = generator.generate_dataset() \n",
+    "vocab = generator.vocab\n",
+    "VOCAB_SIZE = len(vocab)\n",
+    "text_processor = generator.get_text_processor()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Возьмем всопомгательную функцию из пролшло нотубка"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from sklearn.metrics import accuracy_score, precision_recall_fscore_support\n",
+    "\n",
+    "def evaluate_on_test(model, test_loader, device, criterion):\n",
+    "    model.eval()\n",
+    "    total_test_loss = 0\n",
+    "    all_preds = []\n",
+    "    all_labels = []\n",
+    "\n",
+    "    with torch.no_grad():\n",
+    "        for batch_X, batch_y in test_loader:\n",
+    "            batch_X, batch_y = batch_X.to(device), batch_y.to(device)\n",
+    "            outputs = model(batch_X)\n",
+    "            loss = criterion(outputs, batch_y)\n",
+    "            total_test_loss += loss.item()\n",
+    "            \n",
+    "            _, predicted = torch.max(outputs.data, 1)\n",
+    "            all_preds.extend(predicted.cpu().numpy())\n",
+    "            all_labels.extend(batch_y.cpu().numpy())\n",
+    "            \n",
+    "    avg_test_loss = total_test_loss / len(test_loader)\n",
+    "        \n",
+    "    accuracy = accuracy_score(all_labels, all_preds)\n",
+    "    precision, recall, f1, _ = precision_recall_fscore_support(all_labels, all_preds, average='binary')\n",
+    "    \n",
+    "    return {'loss': avg_test_loss, 'accuracy': accuracy, 'precision': precision, 'recall': recall, 'f1_score': f1}"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Создадим генератор датасета и передадим в него уже готовый текстовый процессор, заберем датасет из другого распределения"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "def create_dataloader(X, y, batch_size, shuffle=True):\n",
+    "    X_tensor = torch.as_tensor(X, dtype=torch.long)\n",
+    "    y_tensor = torch.as_tensor(y, dtype=torch.long)\n",
+    "    dataset = TensorDataset(X_tensor, y_tensor)\n",
+    "    return DataLoader(dataset, batch_size=batch_size, shuffle=shuffle)\n",
+    "\n",
+    "text_processor = generator.get_text_processor()\n",
+    "config_polarity = DatasetConfig(\n",
+    "    load_from_disk=True,\n",
+    "    path_to_data=\"../datasets\",\n",
+    "    train_size=25000,  # взяли весь датасет\n",
+    "    val_size=12500,\n",
+    "    test_size=12500,\n",
+    "    build_vocab=False\n",
+    ")\n",
+    "generator_polarity = DatasetGenerator(DatasetName.POLARITY, config=config_polarity)\n",
+    "generator_polarity.vocab = generator.vocab\n",
+    "generator_polarity.id2word = generator.id2word\n",
+    "generator_polarity.text_processor = text_processor\n",
+    "(X_train, y_train), (X_val, y_val), (X_test, y_test) = generator_polarity.generate_dataset()\n",
+    "\n",
+    "\n",
+    "test_loader = create_dataloader(X_test, y_test, BATCH_SIZE, shuffle=False)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Восстановим конфигурации конфигов моделей из прошлого нотубка"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model_configs = {\n",
+    "    \"CustomMamba\": {\n",
+    "        \"class\": CustomMambaClassifier,\n",
+    "        \"params\": {'vocab_size': VOCAB_SIZE, 'd_model': EMBEDDING_DIM, 'd_state': 8, \n",
+    "                   'd_conv': 4, 'num_layers': 2, 'num_classes': NUM_CLASSES},\n",
+    "    },\n",
+    "    \"Lib_LSTM\": {\n",
+    "        \"class\": LSTMClassifier,\n",
+    "        \"params\": {'vocab_size': VOCAB_SIZE, 'embed_dim': EMBEDDING_DIM, 'hidden_dim': 128, \n",
+    "                   'num_layers': 2, 'num_classes': NUM_CLASSES, 'dropout': 0.5},\n",
+    "    },\n",
+    "    \"Lib_Transformer\": {\n",
+    "        \"class\": TransformerClassifier,\n",
+    "        \"params\": {'vocab_size': VOCAB_SIZE, 'embed_dim': EMBEDDING_DIM, 'num_heads': 4, \n",
+    "                   'num_layers': 2, 'num_classes': NUM_CLASSES, 'max_seq_len': MAX_SEQ_LEN},\n",
+    "    },\n",
+    "}"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Теперь посмотрим на результаты"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/gab1k/.cache/pypoetry/virtualenvs/monkey-coding-dl-project-F4QJzkF_-py3.12/lib/python3.12/site-packages/torch/nn/modules/transformer.py:505: UserWarning: The PyTorch API of nested tensors is in prototype stage and will change in the near future. We recommend specifying layout=torch.jagged when constructing a nested tensor, as this layout receives active development, has better operator coverage, and works with torch.compile. (Triggered internally at /pytorch/aten/src/ATen/NestedTensorImpl.cpp:178.)\n",
+      "  output = torch._nested_tensor_from_mask(\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\n",
+      "--- Итоговая таблица сравнения моделей на тестовых данных ---\n",
+      "                     loss  accuracy  precision    recall  f1_score\n",
+      "CustomMamba      0.583675   0.70344   0.653410  0.871734  0.746945\n",
+      "Lib_LSTM         0.675894   0.59520   0.574803  0.744423  0.648709\n",
+      "Lib_Transformer  0.618924   0.66432   0.612190  0.904238  0.730091\n"
+     ]
+    }
+   ],
+   "source": [
+    "results = {}\n",
+    "for model_name, config in model_configs.items():        \n",
+    "    model_path = os.path.join(SAVE_DIR, f\"best_model_{model_name.lower()}.pth\")            \n",
+    "    model = config['class'](**config['params']).to(DEVICE)\n",
+    "\n",
+    "    model.load_state_dict(torch.load(model_path, map_location=DEVICE))\n",
+    "    criterion = nn.CrossEntropyLoss()\n",
+    "    test_metrics = evaluate_on_test(model, test_loader, DEVICE, criterion)\n",
+    "    results[model_name] = test_metrics\n",
+    "    \n",
+    "results_df = pd.DataFrame(results).T\n",
+    "print(\"\\n\\n--- Итоговая таблица сравнения моделей на тестовых данных ---\")\n",
+    "print(results_df.to_string())\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Снимали тут на \"игрушечных данных\". На даже на них видно, что:\n",
+    " - accuracy выше всего на Mamba\n",
+    " - Трансформер справился тоже неплохо\n",
+    " - LSTM опять проиграл\n",
+    "\n",
+    "В следующем нотбуке обучим Mamba и Transformer на всем датасете и снимем качество на втором. Та модель, которая будет лучше, \"поедет в продакшн\" "
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "monkey-coding-dl-project-F4QJzkF_-py3.12",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

notebooks/train.ipynb

@@ -1,68 +1,73 @@
 {
  "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Инициализация глобальных переменных, достаем датасет"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
-    "import warnings\n",
-    "for warn in [UserWarning, FutureWarning]: warnings.filterwarnings(\"ignore\", category = warn)\n",
-    "\n",
     "import os\n",
     "import time\n",
     "import json\n",
     "import torch\n",
+    "import warnings\n",
+    "import numpy as np\n",
+    "import pandas as pd\n",
     "import torch.nn as nn\n",
     "import torch.optim as optim\n",
-    "\n",
     "from torch.utils.data import DataLoader, TensorDataset\n",
+    "from sklearn.metrics import accuracy_score, precision_recall_fscore_support\n",
+    "for warn in [UserWarning, FutureWarning]: warnings.filterwarnings(\"ignore\", category = warn)\n",
     "\n",
-    "# Импортируем классы моделей из нашего файла\n",
-    "from src.models.models import TransformerClassifier, MambaClassifier, LSTMClassifier\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 2,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "MODEL_TO_TRAIN = 'Transformer' \n",
+    "from src.data_utils.config import DatasetConfig\n",
+    "from src.data_utils.dataset_params import DatasetName\n",
+    "from src.data_utils.dataset_generator import DatasetGenerator\n",
+    "from src.models.models import TransformerClassifier\n",
     "\n",
-    "# Гиперпараметры данных и модели\n",
     "MAX_SEQ_LEN = 300\n",
-    "EMBEDDING_DIM = 128\n",
-    "BATCH_SIZE = 32\n",
-    "LEARNING_RATE = 1e-4\n",
-    "NUM_EPOCHS = 5 # Увеличим для лучшего результата\n",
+    "EMBEDDING_DIM = 64 # уменьшили: 128 -> 64, чтобы влезло в гит\n",
+    "BATCH_SIZE = 64 # подняли batch_size: 32 -> 64\n",
+    "LEARNING_RATE = 7e-5\n",
+    "NUM_EPOCHS = 100 # подняли количество эпох: 20 -> 100\n",
+    "NUM_CLASSES = 2\n",
     "\n",
-    "# Пути для сохранения артефактов\n",
     "SAVE_DIR = \"../pretrained\"\n",
     "os.makedirs(SAVE_DIR, exist_ok=True)\n",
     "MODEL_SAVE_PATH = os.path.join(SAVE_DIR, \"best_model.pth\")\n",
     "VOCAB_SAVE_PATH = os.path.join(SAVE_DIR, \"vocab.json\")\n",
     "CONFIG_SAVE_PATH = os.path.join(SAVE_DIR, \"config.json\")\n",
-    "DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n"
+    "DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+    "MODEL_TO_TRAIN = 'Transformer' \n",
+    "\n",
+    "config = DatasetConfig(\n",
+    "    load_from_disk=True,\n",
+    "    path_to_data=\"../datasets\",\n",
+    "    train_size=25000,  # взяли весь датасет\n",
+    "    val_size=12500,\n",
+    "    test_size=12500\n",
+    ")\n",
+    "generator = DatasetGenerator(DatasetName.IMDB, config=config)\n",
+    "(X_train, y_train), (X_val, y_val), (X_test, y_test) = generator.generate_dataset()\n",
+    "VOCAB_SIZE = len(generator.vocab)"
    ]
   },
   {
-   "cell_type": "code",
-   "execution_count": 3,
+   "cell_type": "markdown",
    "metadata": {},
-   "outputs": [],
    "source": [
-    "from src.data_utils.dataset_generator import DatasetGenerator\n",
-    "from src.data_utils.dataset_params import DatasetName\n",
-    "\n",
-    "generator = DatasetGenerator(DatasetName.IMDB)\n",
-    "(X_train, y_train), (X_val, y_val), (X_test, y_test) = generator.generate_dataset()\n",
-    "X_train, y_train, X_val, y_val, X_test, y_test = X_train[:1000], y_train[:1000], X_val[:100], y_val[:100], X_test[:100], y_test[:100]"
+    "Создаем даталоадеры"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -73,45 +78,156 @@
     "val_loader = create_dataloader(X_val, y_val, BATCH_SIZE)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Инициализация модели"
+   ]
+  },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [],
    "source": [
-    "model_params = {}\n",
-    "if MODEL_TO_TRAIN == 'Transformer':\n",
-    "    model_params = {'vocab_size': len(generator.vocab), 'embed_dim': EMBEDDING_DIM, 'num_heads': 4, 'num_layers': 2, 'num_classes': 2, 'max_seq_len': MAX_SEQ_LEN}\n",
-    "    model = TransformerClassifier(**model_params)\n",
-    "elif MODEL_TO_TRAIN == 'Mamba':\n",
-    "    model_params = {'vocab_size': len(generator.vocab), 'embed_dim': EMBEDDING_DIM, 'mamba_d_state': 16, 'mamba_d_conv': 4, 'mamba_expand': 2, 'num_classes': 2}\n",
-    "    model = MambaClassifier(**model_params)\n",
-    "elif MODEL_TO_TRAIN == 'LSTM':\n",
-    "    model_params = {'vocab_size': len(generator.vocab), 'embed_dim': EMBEDDING_DIM, 'hidden_dim': 256, 'num_layers': 2, 'num_classes': 2, 'dropout': 0.5}\n",
-    "    model = LSTMClassifier(**model_params)\n",
-    "else:\n",
-    "    raise ValueError(\"Неизвестный тип модели. Выберите 'Transformer', 'Mamba' или 'LSTM'\")"
+    "model_params = {'vocab_size': len(generator.vocab), 'embed_dim': EMBEDDING_DIM, 'num_heads': 8, 'num_layers': 4, 'num_classes': 2, 'max_seq_len': MAX_SEQ_LEN}\n",
+    "model = TransformerClassifier(**model_params)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Обучение"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "--- Начало обучения модели: Transformer ---\n",
-      "Эпоха 1/5 | Время: 17.06с | Train Loss: 0.7023 | Val Loss: 0.7095 | Val Acc: 0.4000\n",
-      "  -> Модель сохранена, новая лучшая Val Loss: 0.7095\n",
-      "Эпоха 2/5 | Время: 16.40с | Train Loss: 0.6682 | Val Loss: 0.6937 | Val Acc: 0.4800\n",
-      "  -> Модель сохранена, новая лучшая Val Loss: 0.6937\n",
-      "Эпоха 3/5 | Время: 16.13с | Train Loss: 0.6471 | Val Loss: 0.7075 | Val Acc: 0.4100\n",
-      "Эпоха 4/5 | Время: 16.36с | Train Loss: 0.6283 | Val Loss: 0.6917 | Val Acc: 0.5300\n",
-      "  -> Модель сохранена, новая лучшая Val Loss: 0.6917\n",
-      "Эпоха 5/5 | Время: 16.39с | Train Loss: 0.6050 | Val Loss: 0.6871 | Val Acc: 0.5300\n",
-      "  -> Модель сохранена, новая лучшая Val Loss: 0.6871\n"
+      "--- Начало обучения модели ---\n",
+      "Эпоха 1/100 | Время: 14.62с | Train Loss: 0.6563 | Val Loss: 0.6460 | Val Acc: 0.6544\n",
+      "  -> Модель сохранена, новая лучшая Val Loss: 0.6460\n",
+      "Эпоха 2/100 | Время: 14.10с | Train Loss: 0.5749 | Val Loss: 0.5673 | Val Acc: 0.7217\n",
+      "  -> Модель сохранена, новая лучшая Val Loss: 0.5673\n",
+      "Эпоха 3/100 | Время: 14.13с | Train Loss: 0.5058 | Val Loss: 0.5285 | Val Acc: 0.7533\n",
+      "  -> Модель сохранена, новая лучшая Val Loss: 0.5285\n",
+      "Эпоха 4/100 | Время: 14.09с | Train Loss: 0.4664 | Val Loss: 0.4980 | Val Acc: 0.7724\n",
+      "  -> Модель сохранена, новая лучшая Val Loss: 0.4980\n",
+      "Эпоха 5/100 | Время: 14.22с | Train Loss: 0.4382 | Val Loss: 0.4785 | Val Acc: 0.7851\n",
+      "  -> Модель сохранена, новая лучшая Val Loss: 0.4785\n",
+      "Эпоха 6/100 | Время: 14.29с | Train Loss: 0.4166 | Val Loss: 0.4775 | Val Acc: 0.7814\n",
+      "  -> Модель сохранена, новая лучшая Val Loss: 0.4775\n",
+      "Эпоха 7/100 | Время: 14.14с | Train Loss: 0.3974 | Val Loss: 0.4636 | Val Acc: 0.7893\n",
+      "  -> Модель сохранена, новая лучшая Val Loss: 0.4636\n",
+      "Эпоха 8/100 | Время: 14.11с | Train Loss: 0.3778 | Val Loss: 0.4689 | Val Acc: 0.7874\n",
+      "Эпоха 9/100 | Время: 14.30с | Train Loss: 0.3595 | Val Loss: 0.4491 | Val Acc: 0.7973\n",
+      "  -> Модель сохранена, новая лучшая Val Loss: 0.4491\n",
+      "Эпоха 10/100 | Время: 14.23с | Train Loss: 0.3438 | Val Loss: 0.4236 | Val Acc: 0.8148\n",
+      "  -> Модель сохранена, новая лучшая Val Loss: 0.4236\n",
+      "Эпоха 11/100 | Время: 14.45с | Train Loss: 0.3301 | Val Loss: 0.4173 | Val Acc: 0.8174\n",
+      "  -> Модель сохранена, новая лучшая Val Loss: 0.4173\n",
+      "Эпоха 12/100 | Время: 14.21с | Train Loss: 0.3202 | Val Loss: 0.4140 | Val Acc: 0.8206\n",
+      "  -> Модель сохранена, новая лучшая Val Loss: 0.4140\n",
+      "Эпоха 13/100 | Время: 14.10с | Train Loss: 0.3076 | Val Loss: 0.4079 | Val Acc: 0.8243\n",
+      "  -> Модель сохранена, новая лучшая Val Loss: 0.4079\n",
+      "Эпоха 14/100 | Время: 14.07с | Train Loss: 0.2959 | Val Loss: 0.4091 | Val Acc: 0.8220\n",
+      "Эпоха 15/100 | Время: 14.06с | Train Loss: 0.2875 | Val Loss: 0.4074 | Val Acc: 0.8256\n",
+      "  -> Модель сохранена, новая лучшая Val Loss: 0.4074\n",
+      "Эпоха 16/100 | Время: 14.25с | Train Loss: 0.2758 | Val Loss: 0.4021 | Val Acc: 0.8285\n",
+      "  -> Модель сохранена, новая лучшая Val Loss: 0.4021\n",
+      "Эпоха 17/100 | Время: 14.17с | Train Loss: 0.2658 | Val Loss: 0.3933 | Val Acc: 0.8314\n",
+      "  -> Модель сохранена, новая лучшая Val Loss: 0.3933\n",
+      "Эпоха 18/100 | Время: 14.21с | Train Loss: 0.2558 | Val Loss: 0.4100 | Val Acc: 0.8232\n",
+      "Эпоха 19/100 | Время: 14.14с | Train Loss: 0.2518 | Val Loss: 0.3940 | Val Acc: 0.8324\n",
+      "Эпоха 20/100 | Время: 14.14с | Train Loss: 0.2365 | Val Loss: 0.3934 | Val Acc: 0.8304\n",
+      "Эпоха 21/100 | Время: 14.08с | Train Loss: 0.2283 | Val Loss: 0.3913 | Val Acc: 0.8336\n",
+      "  -> Модель сохранена, новая лучшая Val Loss: 0.3913\n",
+      "Эпоха 22/100 | Время: 14.06с | Train Loss: 0.2215 | Val Loss: 0.4161 | Val Acc: 0.8250\n",
+      "Эпоха 23/100 | Время: 14.25с | Train Loss: 0.2100 | Val Loss: 0.3956 | Val Acc: 0.8334\n",
+      "Эпоха 24/100 | Время: 14.25с | Train Loss: 0.2018 | Val Loss: 0.3957 | Val Acc: 0.8333\n",
+      "Эпоха 25/100 | Время: 14.18с | Train Loss: 0.1941 | Val Loss: 0.3942 | Val Acc: 0.8352\n",
+      "Эпоха 26/100 | Время: 14.34с | Train Loss: 0.1811 | Val Loss: 0.3998 | Val Acc: 0.8349\n",
+      "Эпоха 27/100 | Время: 14.18с | Train Loss: 0.1797 | Val Loss: 0.4078 | Val Acc: 0.8318\n",
+      "Эпоха 28/100 | Время: 14.12с | Train Loss: 0.1667 | Val Loss: 0.4101 | Val Acc: 0.8339\n",
+      "Эпоха 29/100 | Время: 14.20с | Train Loss: 0.1610 | Val Loss: 0.4119 | Val Acc: 0.8335\n",
+      "Эпоха 30/100 | Время: 14.25с | Train Loss: 0.1507 | Val Loss: 0.4397 | Val Acc: 0.8294\n",
+      "Эпоха 31/100 | Время: 14.19с | Train Loss: 0.1400 | Val Loss: 0.4245 | Val Acc: 0.8330\n",
+      "Эпоха 32/100 | Время: 14.11с | Train Loss: 0.1361 | Val Loss: 0.4271 | Val Acc: 0.8338\n",
+      "Эпоха 33/100 | Время: 14.13с | Train Loss: 0.1254 | Val Loss: 0.4434 | Val Acc: 0.8311\n",
+      "Эпоха 34/100 | Время: 14.15с | Train Loss: 0.1190 | Val Loss: 0.4446 | Val Acc: 0.8306\n",
+      "Эпоха 35/100 | Время: 14.16с | Train Loss: 0.1125 | Val Loss: 0.4728 | Val Acc: 0.8271\n",
+      "Эпоха 36/100 | Время: 14.23с | Train Loss: 0.1068 | Val Loss: 0.4670 | Val Acc: 0.8297\n",
+      "Эпоха 37/100 | Время: 14.19с | Train Loss: 0.0976 | Val Loss: 0.5572 | Val Acc: 0.8121\n",
+      "Эпоха 38/100 | Время: 14.18с | Train Loss: 0.0923 | Val Loss: 0.4865 | Val Acc: 0.8260\n",
+      "Эпоха 39/100 | Время: 14.16с | Train Loss: 0.0865 | Val Loss: 0.5013 | Val Acc: 0.8250\n",
+      "Эпоха 40/100 | Время: 14.11с | Train Loss: 0.0822 | Val Loss: 0.5205 | Val Acc: 0.8233\n",
+      "Эпоха 41/100 | Время: 14.11с | Train Loss: 0.0731 | Val Loss: 0.5203 | Val Acc: 0.8255\n",
+      "Эпоха 42/100 | Время: 14.17с | Train Loss: 0.0693 | Val Loss: 0.5313 | Val Acc: 0.8276\n",
+      "Эпоха 43/100 | Время: 14.27с | Train Loss: 0.0645 | Val Loss: 0.5518 | Val Acc: 0.8270\n",
+      "Эпоха 44/100 | Время: 14.15с | Train Loss: 0.0577 | Val Loss: 0.5554 | Val Acc: 0.8270\n",
+      "Эпоха 45/100 | Время: 14.23с | Train Loss: 0.0561 | Val Loss: 0.5659 | Val Acc: 0.8258\n",
+      "Эпоха 46/100 | Время: 14.17с | Train Loss: 0.0526 | Val Loss: 0.5840 | Val Acc: 0.8219\n",
+      "Эпоха 47/100 | Время: 14.14с | Train Loss: 0.0457 | Val Loss: 0.6217 | Val Acc: 0.8224\n",
+      "Эпоха 48/100 | Время: 14.27с | Train Loss: 0.0420 | Val Loss: 0.6294 | Val Acc: 0.8237\n",
+      "Эпоха 49/100 | Время: 14.21с | Train Loss: 0.0411 | Val Loss: 0.6333 | Val Acc: 0.8214\n",
+      "Эпоха 50/100 | Время: 14.14с | Train Loss: 0.0345 | Val Loss: 0.6566 | Val Acc: 0.8266\n",
+      "Эпоха 51/100 | Время: 14.12с | Train Loss: 0.0373 | Val Loss: 0.6504 | Val Acc: 0.8243\n",
+      "Эпоха 52/100 | Время: 14.12с | Train Loss: 0.0319 | Val Loss: 0.6640 | Val Acc: 0.8272\n",
+      "Эпоха 53/100 | Время: 14.13с | Train Loss: 0.0286 | Val Loss: 0.6896 | Val Acc: 0.8249\n",
+      "Эпоха 54/100 | Время: 14.14с | Train Loss: 0.0274 | Val Loss: 0.7036 | Val Acc: 0.8213\n",
+      "Эпоха 55/100 | Время: 14.23с | Train Loss: 0.0268 | Val Loss: 0.8750 | Val Acc: 0.7955\n",
+      "Эпоха 56/100 | Время: 14.05с | Train Loss: 0.0274 | Val Loss: 0.7306 | Val Acc: 0.8194\n",
+      "Эпоха 57/100 | Время: 14.06с | Train Loss: 0.0224 | Val Loss: 0.7345 | Val Acc: 0.8196\n",
+      "Эпоха 58/100 | Время: 14.06с | Train Loss: 0.0234 | Val Loss: 0.7029 | Val Acc: 0.8238\n",
+      "Эпоха 59/100 | Время: 14.04с | Train Loss: 0.0218 | Val Loss: 0.7278 | Val Acc: 0.8253\n",
+      "Эпоха 60/100 | Время: 14.15с | Train Loss: 0.0193 | Val Loss: 0.7509 | Val Acc: 0.8217\n",
+      "Эпоха 61/100 | Время: 14.27с | Train Loss: 0.0169 | Val Loss: 0.7706 | Val Acc: 0.8229\n",
+      "Эпоха 62/100 | Время: 14.12с | Train Loss: 0.0177 | Val Loss: 0.7659 | Val Acc: 0.8229\n",
+      "Эпоха 63/100 | Время: 14.35с | Train Loss: 0.0159 | Val Loss: 0.7892 | Val Acc: 0.8178\n",
+      "Эпоха 64/100 | Время: 14.17с | Train Loss: 0.0153 | Val Loss: 0.7721 | Val Acc: 0.8262\n",
+      "Эпоха 65/100 | Время: 14.13с | Train Loss: 0.0161 | Val Loss: 0.7746 | Val Acc: 0.8218\n",
+      "Эпоха 66/100 | Время: 14.14с | Train Loss: 0.0151 | Val Loss: 0.7781 | Val Acc: 0.8227\n",
+      "Эпоха 67/100 | Время: 14.25с | Train Loss: 0.0131 | Val Loss: 0.8032 | Val Acc: 0.8198\n",
+      "Эпоха 68/100 | Время: 14.10с | Train Loss: 0.0156 | Val Loss: 0.7780 | Val Acc: 0.8274\n",
+      "Эпоха 69/100 | Время: 14.04с | Train Loss: 0.0147 | Val Loss: 0.7967 | Val Acc: 0.8237\n",
+      "Эпоха 70/100 | Время: 14.05с | Train Loss: 0.0152 | Val Loss: 0.7833 | Val Acc: 0.8240\n",
+      "Эпоха 71/100 | Время: 14.35с | Train Loss: 0.0136 | Val Loss: 0.8180 | Val Acc: 0.8212\n",
+      "Эпоха 72/100 | Время: 14.24с | Train Loss: 0.0120 | Val Loss: 0.8000 | Val Acc: 0.8235\n",
+      "Эпоха 73/100 | Время: 14.18с | Train Loss: 0.0120 | Val Loss: 0.7985 | Val Acc: 0.8226\n",
+      "Эпоха 74/100 | Время: 14.21с | Train Loss: 0.0106 | Val Loss: 0.7959 | Val Acc: 0.8259\n",
+      "Эпоха 75/100 | Время: 14.12с | Train Loss: 0.0112 | Val Loss: 0.7925 | Val Acc: 0.8238\n",
+      "Эпоха 76/100 | Время: 14.09с | Train Loss: 0.0133 | Val Loss: 0.8455 | Val Acc: 0.8138\n",
+      "Эпоха 77/100 | Время: 14.12с | Train Loss: 0.0099 | Val Loss: 0.8086 | Val Acc: 0.8243\n",
+      "Эпоха 78/100 | Время: 14.19с | Train Loss: 0.0086 | Val Loss: 0.8051 | Val Acc: 0.8271\n",
+      "Эпоха 79/100 | Время: 14.11с | Train Loss: 0.0091 | Val Loss: 0.8212 | Val Acc: 0.8242\n",
+      "Эпоха 80/100 | Время: 14.25с | Train Loss: 0.0105 | Val Loss: 0.8192 | Val Acc: 0.8244\n",
+      "Эпоха 81/100 | Время: 14.20с | Train Loss: 0.0111 | Val Loss: 0.7825 | Val Acc: 0.8250\n",
+      "Эпоха 82/100 | Время: 14.20с | Train Loss: 0.0105 | Val Loss: 0.7885 | Val Acc: 0.8259\n",
+      "Эпоха 83/100 | Время: 14.16с | Train Loss: 0.0091 | Val Loss: 0.7950 | Val Acc: 0.8280\n",
+      "Эпоха 84/100 | Время: 14.27с | Train Loss: 0.0092 | Val Loss: 0.8490 | Val Acc: 0.8217\n",
+      "Эпоха 85/100 | Время: 14.63с | Train Loss: 0.0068 | Val Loss: 0.8464 | Val Acc: 0.8239\n",
+      "Эпоха 86/100 | Время: 14.43с | Train Loss: 0.0084 | Val Loss: 0.8344 | Val Acc: 0.8250\n",
+      "Эпоха 87/100 | Время: 14.27с | Train Loss: 0.0080 | Val Loss: 0.8242 | Val Acc: 0.8266\n",
+      "Эпоха 88/100 | Время: 14.22с | Train Loss: 0.0102 | Val Loss: 0.8427 | Val Acc: 0.8230\n",
+      "Эпоха 89/100 | Время: 14.19с | Train Loss: 0.0080 | Val Loss: 0.8097 | Val Acc: 0.8241\n",
+      "Эпоха 90/100 | Время: 14.24с | Train Loss: 0.0079 | Val Loss: 0.8986 | Val Acc: 0.8161\n",
+      "Эпоха 91/100 | Время: 14.18с | Train Loss: 0.0083 | Val Loss: 0.9104 | Val Acc: 0.8162\n",
+      "Эпоха 92/100 | Время: 14.25с | Train Loss: 0.0073 | Val Loss: 0.8569 | Val Acc: 0.8258\n",
+      "Эпоха 93/100 | Время: 14.17с | Train Loss: 0.0078 | Val Loss: 0.9992 | Val Acc: 0.8039\n",
+      "Эпоха 94/100 | Время: 14.21с | Train Loss: 0.0066 | Val Loss: 0.8613 | Val Acc: 0.8224\n",
+      "Эпоха 95/100 | Время: 14.37с | Train Loss: 0.0067 | Val Loss: 0.8378 | Val Acc: 0.8284\n",
+      "Эпоха 96/100 | Время: 14.24с | Train Loss: 0.0059 | Val Loss: 0.8703 | Val Acc: 0.8203\n",
+      "Эпоха 97/100 | Время: 14.45с | Train Loss: 0.0089 | Val Loss: 0.8341 | Val Acc: 0.8256\n",
+      "Эпоха 98/100 | Время: 14.49с | Train Loss: 0.0131 | Val Loss: 0.8256 | Val Acc: 0.8217\n",
+      "Эпоха 99/100 | Время: 14.46с | Train Loss: 0.0056 | Val Loss: 0.8518 | Val Acc: 0.8202\n",
+      "Эпоха 100/100 | Время: 14.57с | Train Loss: 0.0065 | Val Loss: 0.8770 | Val Acc: 0.8216\n"
      ]
     }
    ],
@@ -121,7 +237,7 @@
     "criterion = nn.CrossEntropyLoss()\n",
     "\n",
     "best_val_loss = float('inf')\n",
-    "print(f\"--- Начало обучения модели: {MODEL_TO_TRAIN} ---\")\n",
+    "print(f\"--- Начало обучения модели ---\")\n",
     "for epoch in range(NUM_EPOCHS):\n",
     "    model.train()\n",
     "    start_time = time.time()\n",
@@ -160,10 +276,143 @@
     "        print(f\"  -> Модель сохранена, новая лучшая Val Loss: {best_val_loss:.4f}\")"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Снимем качество на тестовых данных из исходного датасета"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Метрики на тестовой выборке (из обучаемого датасета) итоговой модели\n",
+      "{'loss': 0.8565363820110049, 'accuracy': 0.8276, 'precision': 0.851743686651778, 'recall': 0.79328, 'f1_score': 0.8214729517024273}\n"
+     ]
+    }
+   ],
+   "source": [
+    "def evaluate_on_test(model, test_loader, device, criterion):\n",
+    "    model.eval()\n",
+    "    total_test_loss = 0\n",
+    "    all_preds = []\n",
+    "    all_labels = []\n",
+    "\n",
+    "    with torch.no_grad():\n",
+    "        for batch_X, batch_y in test_loader:\n",
+    "            batch_X, batch_y = batch_X.to(device), batch_y.to(device)\n",
+    "            outputs = model(batch_X)\n",
+    "            loss = criterion(outputs, batch_y)\n",
+    "            total_test_loss += loss.item()\n",
+    "            \n",
+    "            _, predicted = torch.max(outputs.data, 1)\n",
+    "            all_preds.extend(predicted.cpu().numpy())\n",
+    "            all_labels.extend(batch_y.cpu().numpy())\n",
+    "            \n",
+    "    avg_test_loss = total_test_loss / len(test_loader)\n",
+    "        \n",
+    "    accuracy = accuracy_score(all_labels, all_preds)\n",
+    "    precision, recall, f1, _ = precision_recall_fscore_support(all_labels, all_preds, average='binary')\n",
+    "    \n",
+    "    return {'loss': avg_test_loss, 'accuracy': accuracy, 'precision': precision, 'recall': recall, 'f1_score': f1}\n",
+    "\n",
+    "\n",
+    "test_loader = create_dataloader(X_test, y_test, BATCH_SIZE)\n",
+    "test_metrics = evaluate_on_test(model, test_loader, DEVICE, criterion)\n",
+    "print(f\"Метрики на тестовой выборке (из обучаемого датасета) итоговой модели\\n{test_metrics}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Снимем качество на тестовых данных нового датасета. Считаем данные"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "text_processor = generator.get_text_processor()\n",
+    "config_polarity = DatasetConfig(\n",
+    "    load_from_disk=True,\n",
+    "    path_to_data=\"../datasets\",\n",
+    "    train_size=25000,  # взяли весь датасет\n",
+    "    val_size=12500,\n",
+    "    test_size=12500,\n",
+    "    build_vocab=False\n",
+    ")\n",
+    "generator_polarity = DatasetGenerator(DatasetName.POLARITY, config=config_polarity)\n",
+    "generator_polarity.vocab = generator.vocab\n",
+    "generator_polarity.id2word = generator.id2word\n",
+    "generator_polarity.text_processor = text_processor\n",
+    "(X_train, y_train), (X_val, y_val), (X_test, y_test) = generator_polarity.generate_dataset()\n"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 7,
    "metadata": {},
+   "outputs": [],
+   "source": [
+    "test_loader = create_dataloader(X_test, y_test, BATCH_SIZE)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Посмтрим на метрики"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Метрики на тестовой выборке (из неизвестного датасета) итоговой модели\n",
+      "{'loss': 0.6786724476485836, 'accuracy': 0.73816, 'precision': 0.7227414330218068, 'recall': 0.7762906309751434, 'f1_score': 0.7485595759391565}\n"
+     ]
+    }
+   ],
+   "source": [
+    "test_metrics = evaluate_on_test(model, test_loader, DEVICE, criterion)\n",
+    "print(f\"Метрики на тестовой выборке (из неизвестного датасета) итоговой модели\\n{test_metrics}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "В целом видно, что модель что-то, да выучила. Гипотезы по улучшению:\n",
+    " - Больше и разнообразнее данные для обучения\n",
+    "    - Чем больше словарь - тем лучше\n",
+    " - Нужно чтобы тестовый датасет был больше похож на обучаемый"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Сохранение итоговой модели"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
    "outputs": [
     {
      "name": "stdout",
@@ -184,20 +433,13 @@
     "}\n",
     "with open(CONFIG_SAVE_PATH, 'w', encoding='utf-8') as f:\n",
     "    json.dump(config, f, ensure_ascii=False, indent=4)\n",
-    "print(f\"Конфигурация модели сохранена в: {CONFIG_SAVE_PATH}\")\n"
+    "print(f\"Конфигурация модели сохранена в: {CONFIG_SAVE_PATH}\")"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "monkey-coding-dl-project-OWiM8ypK-py3.12",
+   "display_name": "monkey-coding-dl-project-F4QJzkF_-py3.12",
    "language": "python",
    "name": "python3"
   },
@@ -211,7 +453,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.3"
+   "version": "3.12.11"
   }
  },
  "nbformat": 4,

pretrained/best_model.pth

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:7636cf4c7205b64df4b91b2a23620d443de468a91211e074760e64adb24751ba
-size 37445685
+oid sha256:0709ee349a04faf4ff0e19d0a91967953dd8aa0ca2fb0863f866fcc3d219debc
+size 29525749

pretrained/config.json

@@ -2,10 +2,10 @@
     "model_type": "Transformer",
     "max_seq_len": 300,
     "model_params": {
-        "vocab_size": 69715,
-        "embed_dim": 128,
-        "num_heads": 4,
-        "num_layers": 2,
+        "vocab_size": 111829,
+        "embed_dim": 64,
+        "num_heads": 8,
+        "num_layers": 4,
         "num_classes": 2,
         "max_seq_len": 300
     }

src/data_utils/config.py

@@ -15,6 +15,7 @@ class DatasetConfig:
         min_word_freq: Minimum word frequency to include in vocabulary
         load_from_disk: Load dataset from local dir. If false download from huggin face
         path_to_data: Path to local dataset data
+        build_vocab: Is build vocabulary necessary
         max_seq_len: Maximum sequence length (will be padded/truncated to this)
         lowercase: Whether to convert text to lowercase
         remove_punct: Whether to remove punctuation
@@ -30,6 +31,7 @@ class DatasetConfig:
     min_word_freq: int = 1
     load_from_disk: bool = False
     path_to_data: str = "./datasets"
+    build_vocab: bool = True
 
     max_seq_len: int = 300
     lowercase: bool = True

src/data_utils/dataset_generator.py

@@ -128,7 +128,8 @@ class DatasetGenerator:
         train_texts = train_df[self.dataset_params.content_col_name].tolist()
         train_tokens = [self.text_processor.preprocess_text(text) for text in train_texts]
         
-        self.vocab, self.id2word = self.build_vocabulary(train_tokens)
+        if self.config.build_vocab:
+            self.vocab, self.id2word = self.build_vocabulary(train_tokens)
 
         X_train = torch.stack([self.text_processor.text_to_tensor(text) for text in train_texts])