naijagpt.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "mount_file_id": "1g4vVXDsdVH_oKrUtiv_ucZvdvWLxeNlK",
      "authorship_tag": "ABX9TyM/OAUnlDGSppkk2B6OEx2J",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Deepersensor/NaijaGPT/blob/master/naijagpt.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nQgezvXF_Nbd"
      },
      "outputs": [],
      "source": [
        "# prompt: build a python notebook that attempts to train and build a text conversational model from scratch also using data (.parquet files) in the data/folder\n",
        "\n",
        "# Install necessary libraries\n",
        "!pip install transformers datasets evaluate rouge_score pandas pyarrow\n",
        "\n",
        "# Import libraries\n",
        "import pandas as pd\n",
        "from datasets import Dataset\n",
        "from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, TrainingArguments, Trainer\n",
        "import evaluate\n",
        "import numpy as np\n",
        "\n",
        "# Load data from Parquet files (adjust paths as needed)\n",
        "def load_data_from_parquet(data_dir):\n",
        "  data_files = {\n",
        "      'train': data_dir + 'train.parquet',\n",
        "      'validation': data_dir + 'validation.parquet',\n",
        "      # Add test data if available: 'test': data_dir + 'test.parquet',\n",
        "  }\n",
        "  dataset = Dataset.from_parquet(data_files)\n",
        "  return dataset\n",
        "\n",
        "# Replace 'data/' with the actual path to your data folder\n",
        "dataset = load_data_from_parquet('data/')\n",
        "\n",
        "# Preprocess data (customize based on your dataset structure)\n",
        "def preprocess_function(examples):\n",
        "  model_inputs = tokenizer(examples['input_text'], max_length=1024, truncation=True)\n",
        "  with tokenizer.as_target_tokenizer():\n",
        "    labels = tokenizer(examples['output_text'], max_length=1024, truncation=True)\n",
        "  model_inputs[\"labels\"] = labels[\"input_ids\"]\n",
        "  return model_inputs\n",
        "\n",
        "# Choose a pre-trained model for sequence-to-sequence tasks\n",
        "model_checkpoint = \"google/flan-t5-base\" # Or another model like \"t5-base\"\n",
        "tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)\n",
        "model = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint)\n",
        "\n",
        "tokenized_datasets = dataset.map(preprocess_function, batched=True)\n",
        "\n",
        "# Define evaluation metrics (customize based on your task)\n",
        "metric = evaluate.load(\"rouge\")\n",
        "\n",
        "def compute_metrics(eval_pred):\n",
        "  predictions, labels = eval_pred\n",
        "  decoded_preds = tokenizer.batch_decode(predictions, skip_special_tokens=True)\n",
        "  labels = np.where(labels != -100, labels, tokenizer.pad_token_id)\n",
        "  decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)\n",
        "\n",
        "  result = metric.compute(predictions=decoded_preds, references=decoded_labels, use_stemmer=True)\n",
        "  # Extract ROUGE scores\n",
        "  result = {key: value.mid.fmeasure * 100 for key, value in result.items()}\n",
        "  return {k: round(v, 4) for k, v in result.items()}\n",
        "\n",
        "# Define training arguments\n",
        "training_args = TrainingArguments(\n",
        "    output_dir=\"./results\",\n",
        "    evaluation_strategy=\"epoch\",\n",
        "    learning_rate=2e-5,\n",
        "    per_device_train_batch_size=16,\n",
        "    per_device_eval_batch_size=16,\n",
        "    weight_decay=0.01,\n",
        "    save_total_limit=3,\n",
        "    num_train_epochs=3,\n",
        "    predict_with_generate=True,\n",
        "    fp16=True, # Use mixed precision training if supported by your hardware\n",
        ")\n",
        "\n",
        "# Initialize the Trainer\n",
        "trainer = Trainer(\n",
        "    model=model,\n",
        "    args=training_args,\n",
        "    train_dataset=tokenized_datasets[\"train\"],\n",
        "    eval_dataset=tokenized_datasets[\"validation\"],\n",
        "    tokenizer=tokenizer,\n",
        "    compute_metrics=compute_metrics,\n",
        ")\n",
        "\n",
        "# Train the model\n",
        "trainer.train()\n",
        "\n",
        "# Evaluate the model\n",
        "trainer.evaluate()\n",
        "\n",
        "# Save the trained model\n",
        "trainer.save_model(\"./trained_model\")\n",
        "\n",
        "# Example of how to use the trained model for inference\n",
        "input_text = \"How are you doing?\"\n",
        "inputs = tokenizer(input_text, return_tensors=\"pt\").to(\"cuda\")\n",
        "outputs = model.generate(**inputs)\n",
        "decoded_output = tokenizer.decode(outputs[0], skip_special_tokens=True)\n",
        "decoded_output"
      ]
    }
  ]
}