{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "# Install necessary libraries\n",
        "!pip install datasets\n",
        "\n",
        "# Importing required libraries for dataset, model, tokenizer, training, and evaluation\n",
        "from datasets import load_dataset\n",
        "from transformers import T5Tokenizer, T5ForConditionalGeneration, Trainer, TrainingArguments\n",
        "from sklearn.model_selection import train_test_split\n",
        "import torch\n",
        "from IPython import get_ipython\n",
        "from IPython.display import display"
      ],
      "metadata": {
        "id": "HrXBpUDmLbKH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# 1. Load Dataset\n",
        "# Load the Bengali transliteration dataset\n",
        "raw_dataset = load_dataset(\"SKNahin/bengali-transliteration-data\")\n",
        "\n",
        "# Split dataset into training and validation sets (90% training, 10% validation)\n",
        "train_data = raw_dataset['train'].train_test_split(test_size=0.1, seed=42)['train'] # Added seed for reproducibility\n",
        "val_data = raw_dataset['train'].train_test_split(test_size=0.1, seed=42)['test'] # Added seed for reproducibility\n",
        "\n",
        "# 2. Preprocessing\n",
        "# Define model name for T5 and load its tokenizer\n",
        "model_name = \"google/mt5-small\"\n",
        "tokenizer = T5Tokenizer.from_pretrained(model_name)\n",
        "\n",
        "# Tokenize function for preprocessing the data\n",
        "def preprocess_data(examples):\n",
        "    # Tokenize input and target sequences with padding and truncation to fixed length\n",
        "    # Access the correct columns based on the dataset structure\n",
        "    inputs = tokenizer(examples['bn'], padding=\"max_length\", truncation=True, max_length=128)\n",
        "    targets = tokenizer(examples['rm'], padding=\"max_length\", truncation=True, max_length=128)\n",
        "\n",
        "    # Assign the tokenized target as labels for the model\n",
        "    inputs['labels'] = targets['input_ids']\n",
        "    return inputs\n",
        "\n",
        "# Apply the preprocessing function to the training and validation datasets\n",
        "train_dataset = train_data.map(preprocess_data, batched=True)\n",
        "val_dataset = val_data.map(preprocess_data, batched=True)\n",
        "\n",
        "# 3. Model Selection\n",
        "# Load the T5 model for conditional generation\n",
        "model = T5ForConditionalGeneration.from_pretrained(model_name)\n",
        "\n",
        "# 4. Training\n",
        "# Define training arguments like learning rate, batch size, number of epochs, etc.\n",
        "training_args = TrainingArguments(\n",
        "    output_dir=\"./results\", # Directory to store the results\n",
        "    evaluation_strategy=\"epoch\", # Evaluate model after each epoch\n",
        "    learning_rate=5e-5, # Learning rate\n",
        "    per_device_train_batch_size=16, # Batch size for training\n",
        "    per_device_eval_batch_size=16, # Batch size for evaluation\n",
        "    num_train_epochs=5, # Number of training epochs\n",
        "    weight_decay=0.01, # Weight decay for regularization\n",
        "    save_steps=1000, # Save model every 1000 steps\n",
        "    save_total_limit=2, # Keep only 2 most recent checkpoints\n",
        "    logging_dir=\"./logs\", # Directory to store logs\n",
        "    logging_steps=500, # Log every 500 steps\n",
        ")\n"
      ],
      "metadata": {
        "id": "kxHZMLRKPr6L"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Initialize the Trainer with the model, arguments, and datasets\n",
        "trainer = Trainer(\n",
        "    model=model,\n",
        "    args=training_args,\n",
        "    train_dataset=train_dataset,\n",
        "    eval_dataset=val_dataset\n",
        ")\n",
        "\n",
        "# Start the training process\n",
        "trainer.train()"
      ],
      "metadata": {
        "id": "9_9U0GCCoTqZ"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}