train_pt.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "\n",
    "import pandas as pd\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from sklearn.model_selection import train_test_split\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "from _functions import build_vocab, convert_df_id, plot_result\n",
    "from _objects import Decoder, Encoder, Seq2Seq, TextNormalizationDataset\n",
    "\n",
    "# The seed for random state\n",
    "_SEED: int = 574\n",
    "\n",
    "# Set seed for reproducibility\n",
    "torch.manual_seed(_SEED)\n",
    "\n",
    "# Get the device for training\n",
    "device: torch.device = torch.device(\"cuda\")\n",
    "\n",
    "# Load training data\n",
    "df: pd.DataFrame = convert_df_id(\n",
    "    pd.read_csv(\"en_train.csv.zip\")\n",
    "    .drop(columns=[\"class\"])\n",
    "    .astype(str)\n",
    "    .drop_duplicates(subset=[\"before\", \"after\"], keep=\"first\")\n",
    ")\n",
    "\n",
    "# create train/test split\n",
    "train_df, test_df = train_test_split(df, test_size=0.2, random_state=_SEED)\n",
    "\n",
    "# Vocabularies for input and output\n",
    "input_vocab: dict[str, int] = {}\n",
    "output_vocab: dict[str, int] = {}\n",
    "\n",
    "# Load existing vocabularies for input and output\n",
    "vocab_path: str = \"vocab.json\"\n",
    "if os.path.exists(vocab_path):\n",
    "    with open(vocab_path, encoding=\"utf-8\") as f:\n",
    "        d: dict[str, dict[str, int]] = json.load(f)\n",
    "        input_vocab = d[\"input_vocab\"]\n",
    "        output_vocab = d[\"output_vocab\"]\n",
    "# Build vocabularies for input and output\n",
    "else:\n",
    "    input_vocab = build_vocab(train_df[\"before\"])\n",
    "    output_vocab = build_vocab(train_df[\"after\"])\n",
    "    with open(vocab_path, \"w\", encoding=\"utf-8\") as f:\n",
    "        json.dump({\"input_vocab\": input_vocab, \"output_vocab\": output_vocab}, f)\n",
    "\n",
    "# Inverse vocabularies for decoding\n",
    "inverse_input_vocab: dict[str, int] = {v: k for k, v in input_vocab.items()}\n",
    "inverse_output_vocab: dict[str, int] = {v: k for k, v in output_vocab.items()}\n",
    "\n",
    "\n",
    "# Define model hyperparameters\n",
    "INPUT_DIM: int = 2560  # len(input_vocab)\n",
    "OUTPUT_DIM: int = 2560  # len(output_vocab)\n",
    "ENC_EMB_DIM: int = 256\n",
    "DEC_EMB_DIM: int = 256\n",
    "HID_DIM: int = 512\n",
    "N_LAYERS: int = 2\n",
    "ENC_DROPOUT: float = 0.2\n",
    "DEC_DROPOUT: float = 0.2\n",
    "BATCH_SIZE: int = 128\n",
    "\n",
    "# Create the encoder and decoder\n",
    "encoder: Encoder = Encoder(INPUT_DIM, ENC_EMB_DIM, HID_DIM, N_LAYERS, ENC_DROPOUT)\n",
    "decoder: Decoder = Decoder(OUTPUT_DIM, DEC_EMB_DIM, HID_DIM, N_LAYERS, DEC_DROPOUT)\n",
    "\n",
    "# Create the model\n",
    "model: Seq2Seq = Seq2Seq(encoder, decoder, device).to(device)\n",
    "\n",
    "# Define the loss function and optimizer\n",
    "criterion: nn.CrossEntropyLoss = nn.CrossEntropyLoss(ignore_index=output_vocab[\"<pad>\"])\n",
    "optimizer: optim.Optimizer = optim.Adam(model.parameters())\n",
    "\n",
    "# Prepare the data\n",
    "train_dataset: TextNormalizationDataset = TextNormalizationDataset(\n",
    "    train_df, input_vocab, output_vocab\n",
    ")\n",
    "train_loader: DataLoader = DataLoader(\n",
    "    train_dataset, batch_size=BATCH_SIZE, shuffle=True\n",
    ")\n",
    "\n",
    "# Create a validation set from the training data\n",
    "train_df, val_df = train_test_split(train_df, test_size=0.2, random_state=_SEED)\n",
    "val_dataset: TextNormalizationDataset = TextNormalizationDataset(\n",
    "    val_df, input_vocab, output_vocab\n",
    ")\n",
    "val_loader: DataLoader = DataLoader(val_dataset, batch_size=BATCH_SIZE)\n",
    "\n",
    "# Testing dataset\n",
    "test_dataset: TextNormalizationDataset = TextNormalizationDataset(\n",
    "    test_df, input_vocab, output_vocab\n",
    ")\n",
    "test_loader: DataLoader = DataLoader(test_dataset, batch_size=BATCH_SIZE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to train the model\n",
    "def train(\n",
    "    model: Seq2Seq,\n",
    "    dataloader: DataLoader,\n",
    "    optimizer: optim.Optimizer,\n",
    "    criterion: nn.CrossEntropyLoss,\n",
    ") -> float:\n",
    "    model.train()\n",
    "    epoch_loss: int = 0\n",
    "\n",
    "    # Create a progress bar\n",
    "    total_batches: int = len(dataloader)\n",
    "    print(f\"\\nTraining: 0/{total_batches} batches processed\", end=\"\\r\")\n",
    "\n",
    "    for i, batch in enumerate(dataloader):\n",
    "        src = batch[\"input\"].to(device)\n",
    "        trg = batch[\"output\"].to(device)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        output = model(src, trg)\n",
    "\n",
    "        # trg: [batch_size, trg_len]\n",
    "        # output: [batch_size, trg_len, output_dim]\n",
    "\n",
    "        output_dim = output.shape[-1]\n",
    "\n",
    "        # Exclude the first token (<sos>)\n",
    "        output = output[:, 1:].reshape(-1, output_dim)\n",
    "        trg = trg[:, 1:].reshape(-1)\n",
    "\n",
    "        # Calculate loss\n",
    "        loss = criterion(output, trg)\n",
    "\n",
    "        # Backpropagation\n",
    "        loss.backward()\n",
    "\n",
    "        # Gradient clipping\n",
    "        torch.nn.utils.clip_grad_norm_(model.parameters(), 1)\n",
    "\n",
    "        # Update parameters\n",
    "        optimizer.step()\n",
    "\n",
    "        epoch_loss += loss.item()\n",
    "\n",
    "        # Update progress bar every 10 batches\n",
    "        if (i + 1) % 10 == 0 or (i + 1) == total_batches:\n",
    "            print(\n",
    "                f\"Training: {i+1}/{total_batches} batches processed | Current loss: {loss.item():.5f}\",\n",
    "                end=\"\\r\",\n",
    "            )\n",
    "\n",
    "    avg_loss: float = epoch_loss / len(dataloader)\n",
    "    print(f\"\\nTraining completed. Average batch loss: {avg_loss:.5f}\" + \" \" * 30)\n",
    "    return avg_loss\n",
    "\n",
    "\n",
    "# Function to evaluate the model\n",
    "def evaluate(\n",
    "    model: Seq2Seq, dataloader: DataLoader, criterion: nn.CrossEntropyLoss\n",
    ") -> float:\n",
    "    model.eval()\n",
    "    epoch_loss: int = 0\n",
    "\n",
    "    # Create a progress bar\n",
    "    total_batches: int = len(dataloader)\n",
    "    print(f\"Evaluating: 0/{total_batches} batches processed\", end=\"\\r\")\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for i, batch in enumerate(dataloader):\n",
    "            src = batch[\"input\"].to(device)\n",
    "            trg = batch[\"output\"].to(device)\n",
    "\n",
    "            output = model(src, trg, 0)  # Turn off teacher forcing\n",
    "\n",
    "            # trg: [batch_size, trg_len]\n",
    "            # output: [batch_size, trg_len, output_dim]\n",
    "\n",
    "            output_dim = output.shape[-1]\n",
    "\n",
    "            output = output[:, 1:].reshape(-1, output_dim)\n",
    "            trg = trg[:, 1:].reshape(-1)\n",
    "\n",
    "            loss = criterion(output, trg)\n",
    "\n",
    "            epoch_loss += loss.item()\n",
    "\n",
    "            # Update progress bar periodically\n",
    "            if (i + 1) % 10 == 0 or (i + 1) == total_batches:\n",
    "                print(f\"Evaluating: {i+1}/{total_batches} batches processed\", end=\"\\r\")\n",
    "\n",
    "    avg_loss: float = epoch_loss / len(dataloader)\n",
    "    print(f\"\\nEvaluation completed. Average batch loss: {avg_loss:.5f}\" + \" \" * 30)\n",
    "    return avg_loss\n",
    "\n",
    "\n",
    "# Function to predict on test data\n",
    "def predict(\n",
    "    model: Seq2Seq, dataloader: DataLoader, inverse_output_vocab: dict[str, int]\n",
    "):\n",
    "    model.eval()\n",
    "    predictions = []\n",
    "    ids = []\n",
    "\n",
    "    # Create a progress bar\n",
    "    total_batches: int = len(dataloader)\n",
    "    print(f\"Predicting: 0/{total_batches} batches processed\", end=\"\\r\")\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for i, batch in enumerate(dataloader):\n",
    "            src = batch[\"input\"].to(device)\n",
    "            id_batch = batch[\"id\"]\n",
    "\n",
    "            batch_size = src.shape[0]\n",
    "\n",
    "            # Encoder outputs\n",
    "            _, hidden = model.encoder(src)\n",
    "\n",
    "            # First input to the decoder is the <sos> token\n",
    "            input = torch.tensor([output_vocab[\"<sos>\"]] * batch_size).to(device)\n",
    "\n",
    "            # Store predictions for each sequence in the batch\n",
    "            batch_outputs = [[] for _ in range(batch_size)]\n",
    "\n",
    "            for _ in range(50):  # Adjust max length as needed\n",
    "                output, hidden = model.decoder(input, hidden)\n",
    "                pred_tokens = output.argmax(1)\n",
    "\n",
    "                # Add predicted tokens to each sequence's output\n",
    "                for i, token in enumerate(pred_tokens):\n",
    "                    batch_outputs[i].append(token.item())\n",
    "\n",
    "                # Update input for next time step\n",
    "                input = pred_tokens\n",
    "\n",
    "            # Convert token indices to characters for each sequence\n",
    "            for i in range(batch_size):\n",
    "                tokens = batch_outputs[i]\n",
    "                # Find where the sequence ends (at <eos> token)\n",
    "                if output_vocab[\"<eos>\"] in tokens:\n",
    "                    end_idx = tokens.index(output_vocab[\"<eos>\"])\n",
    "                    tokens = tokens[:end_idx]\n",
    "\n",
    "                # Convert tokens to text\n",
    "                text: str = \"\".join(\n",
    "                    [\n",
    "                        inverse_output_vocab.get(token, \"\")\n",
    "                        for token in tokens\n",
    "                        if token\n",
    "                        not in [\n",
    "                            output_vocab[\"<sos>\"],\n",
    "                            output_vocab[\"<eos>\"],\n",
    "                            output_vocab[\"<pad>\"],\n",
    "                        ]\n",
    "                    ]\n",
    "                )\n",
    "\n",
    "                predictions.append(text)\n",
    "                ids.append(id_batch[i])\n",
    "\n",
    "            # Update progress bar periodically\n",
    "            if (i + 1) % 10 == 0 or (i + 1) == total_batches:\n",
    "                print(f\"Predicting: {i+1}/{total_batches} batches processed\", end=\"\\r\")\n",
    "\n",
    "    print(\n",
    "        f\"\\nPrediction completed. Generated {len(predictions)} predictions.\" + \" \" * 30\n",
    "    )\n",
    "    return ids, predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training loop\n",
    "N_EPOCHS: int = 50\n",
    "# How many epochs for early stopping, -1 to disable, 0 to stop immediately\n",
    "early_stop_patience: int = 5\n",
    "# Rounds since last best loss\n",
    "current_patience: int = 0\n",
    "# Best validation lost\n",
    "best_valid_loss: float = float(\"inf\")\n",
    "\n",
    "print(f\"Starting training for {N_EPOCHS} epochs...\")\n",
    "print(f\"Training on {len(train_dataset)} examples\")\n",
    "print(f\"Validating on {len(val_dataset)} examples\")\n",
    "print(f\"Testing on {len(test_dataset)} examples\")\n",
    "\n",
    "# For tracking metrics\n",
    "train_losses: list[float] = []\n",
    "valid_losses: list[float] = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for epoch in range(N_EPOCHS):\n",
    "\n",
    "    print(f\"\\n{'='*50}\")\n",
    "    print(f\"Epoch: {epoch+1}/{N_EPOCHS}\")\n",
    "\n",
    "    train_loss: float = train(model, train_loader, optimizer, criterion)\n",
    "    valid_loss: float = evaluate(model, val_loader, criterion)\n",
    "\n",
    "    train_losses.append(train_loss)\n",
    "    valid_losses.append(valid_loss)\n",
    "\n",
    "    print(f\"Train Loss: {train_loss:.5f}\")\n",
    "    print(f\"Validation Loss: {valid_loss:.5f}\")\n",
    "\n",
    "    if valid_loss < best_valid_loss:\n",
    "        best_valid_loss = valid_loss\n",
    "        current_patience = 0\n",
    "        torch.save(model.state_dict(), \"best-model.pt\")\n",
    "        print(f\"New best model saved with validation loss: {valid_loss:.5f}\")\n",
    "    else:\n",
    "        current_patience += 1\n",
    "        print(\n",
    "            f\"{current_patience}/{early_stop_patience} epochs since loss decreases: {best_valid_loss:.5f}\"\n",
    "        )\n",
    "\n",
    "    # Early stopping check\n",
    "    if early_stop_patience >= 0 and current_patience > early_stop_patience:\n",
    "        print(\n",
    "            f\"Validation loss increasing for {early_stop_patience} consecutive epochs. Early stopping...\"\n",
    "        )\n",
    "        break\n",
    "\n",
    "print(f\"\\n{'='*50}\")\n",
    "print(\"Training complete!\")\n",
    "print(f\"Best validation loss: {best_valid_loss:.5f}\")\n",
    "\n",
    "# Print training summary\n",
    "print(\"\\nTraining Summary:\")\n",
    "for epoch, (t_loss, v_loss) in enumerate(zip(train_losses, valid_losses), 1):\n",
    "    print(f\"Epoch {epoch}: Train Loss = {t_loss:.5f}, Valid Loss = {v_loss:.5f}\")\n",
    "\n",
    "# Plotting training/validation loss\n",
    "plot_result(train_losses, valid_losses)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load best model for prediction\n",
    "print(\"\\nLoading best model for testing...\")\n",
    "model.load_state_dict(torch.load(\"best-model.pt\"))\n",
    "\n",
    "# Create a dictionary mapping ids to indices for easier lookup\n",
    "test_id_to_idx: dict[str, int] = {\n",
    "    id_val: i for i, id_val in enumerate(test_df[\"id\"].astype(str))\n",
    "}\n",
    "\n",
    "# Make predictions on test data\n",
    "ids, predictions = predict(model, test_loader, inverse_output_vocab)\n",
    "\n",
    "# Prepare result lists\n",
    "actual_values: list[str] = []\n",
    "before_values: list[str] = []\n",
    "\n",
    "# Match predictions with actual values using ids\n",
    "for pred_id in ids:\n",
    "    if pred_id in test_id_to_idx:\n",
    "        idx = test_id_to_idx[pred_id]\n",
    "        actual_values.append(test_df.iloc[idx][\"after\"])\n",
    "        before_values.append(test_df.iloc[idx][\"before\"])\n",
    "    else:\n",
    "        # If id not found, use empty strings as placeholders\n",
    "        actual_values.append(\"\")\n",
    "        before_values.append(\"\")\n",
    "\n",
    "# Calculate accuracy\n",
    "correct: int = sum(1 for act, pred in zip(actual_values, predictions) if act == pred)\n",
    "accuracy: float = correct * 100 / len(actual_values) if actual_values else 0\n",
    "print(f\"Test Accuracy: {accuracy:.2f}%\")\n",
    "\n",
    "# Save predictions for analysis\n",
    "results_df: pd.DataFrame = pd.DataFrame(\n",
    "    {\n",
    "        \"id\": ids,\n",
    "        \"before\": before_values,\n",
    "        \"actual\": actual_values,\n",
    "        \"predicted\": predictions,\n",
    "    }\n",
    ")\n",
    "results_df.to_csv(\"en_test_submission.csv\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load best model for prediction\n",
    "print(\"\\nLoading best model for submission...\")\n",
    "model.load_state_dict(torch.load(\"best-model.pt\"))\n",
    "\n",
    "# Load and process Validation dataframe\n",
    "val_df: pd.DataFrame = pd.read_csv(\"./en_test.csv.zip\")\n",
    "val_df[\"id\"] = val_df[\"sentence_id\"].astype(str) + \"_\" + val_df[\"token_id\"].astype(str)\n",
    "val_df = val_df.drop(columns=[\"sentence_id\", \"token_id\"]).astype(str)\n",
    "\n",
    "# Validation dataset\n",
    "val_dataset: TextNormalizationDataset = TextNormalizationDataset(\n",
    "    val_df, input_vocab, output_vocab\n",
    ")\n",
    "val_loader: DataLoader = DataLoader(val_dataset, batch_size=BATCH_SIZE)\n",
    "\n",
    "# Make predictions on test data\n",
    "ids, predictions = predict(model, val_loader, inverse_output_vocab)\n",
    "\n",
    "# Save predictions for analysis\n",
    "en_sample_submission = pd.DataFrame(\n",
    "    {\n",
    "        \"id\": ids,\n",
    "        \"after\": predictions,\n",
    "    }\n",
    ").astype(str)\n",
    "\n",
    "en_sample_submission.to_csv(\"en_sample_submission.csv\", index=False)\n",
    "# en_sample_submission.to_csv(\"en_sample_submission.csv.zip\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "en_sample_submission = (\n",
    "    pd.read_csv(\"en_sample_submission.csv\").astype(str).set_index(\"id\")\n",
    ")\n",
    "en_sample_submission_true = (\n",
    "    pd.read_csv(\"en_sample_submission.csv.zip\").astype(str).set_index(\"id\")\n",
    ")\n",
    "merged_df = en_sample_submission.merge(\n",
    "    en_sample_submission_true, on=\"id\", suffixes=(\"\", \"_df2\")\n",
    ")\n",
    "\n",
    "# Calculate accuracy\n",
    "correct: int = (merged_df[\"after\"] == merged_df[\"after_df2\"]).sum()\n",
    "accuracy: float = correct * 100 / len(merged_df)\n",
    "print(f\"Test Accuracy: {accuracy:.2f}%\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}