data_preprocessing.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3bb3de21-8909-4200-8da8-340f79252f4b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!pip install datasets\n",
    "!pip install transformers\n",
    "!pip install tqdm\n",
    "!pip install nltk\n",
    "!pip install bitsandbytes\n",
    "!pip install accelerate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b7b4e38-bf80-4d1a-9fc5-bc598d172363",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "from tqdm import tqdm\n",
    "import json\n",
    "import nltk\n",
    "from transformers import AutoTokenizer\n",
    "from joblib import Parallel, delayed\n",
    "import numpy as np\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig\n",
    "import torch\n",
    "import ijson\n",
    "\n",
    "nltk.download('punkt')\n",
    "nltk.download('punkt_tab')\n",
    "model_id = \"NousResearch/Llama-2-7b-hf\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_id, padding_side=\"right\", user_fast=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cfcab10f",
   "metadata": {},
   "source": [
    "# Load dataset and store it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9578caa5-57e1-4ef7-a995-9555663ae68a",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = load_dataset(\"wikipedia\", \"20220301.en\", trust_remote_code=True)\n",
    "training_data = dataset[\"train\"][\"text\"]\n",
    "with open(\"training_data.json\", 'w') as file:\n",
    "    json.dump(training_data, file, indent=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bdcd58e6",
   "metadata": {},
   "source": [
    "# Split into sentences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c598d72-6ad1-4327-9eeb-fc7ab4d2866c",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"training_data.json\", 'r') as file:\n",
    "    training_data = json.load(file)\n",
    "\n",
    "def split_into_sentences(text: str) -> list[str]:\n",
    "    return nltk.sent_tokenize(text)\n",
    "\n",
    "training_data = Parallel(n_jobs=-1)(delayed(split_into_sentences)(webpage) for webpage in tqdm(training_data, desc=\"Processing\"))\n",
    "\n",
    "with open(\"sentence_training_data.json\", 'w') as file:\n",
    "    json.dump(training_data, file, indent=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8813691f",
   "metadata": {},
   "source": [
    "# Tokenize training data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02251c1f-649c-4cd8-985d-54d1c7fd180b",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"sentence_training_data.json\", 'r') as file:\n",
    "    training_data = json.load(file)\n",
    "\n",
    "def tokenize_function(sentences: list[str]) -> list[list[int]]:\n",
    "    for i in range(len(sentences)):\n",
    "        sentences[i] = tokenizer(sentences[i])[\"input_ids\"]\n",
    "    return sentences\n",
    "\n",
    "# The tokenization process cannot be executed in parallel.\n",
    "for i in tqdm(range(len(training_data))):\n",
    "    for j in range(len(training_data[i])):\n",
    "        training_data[i][j] = tokenizer(training_data[i][j])[\"input_ids\"]\n",
    "\n",
    "with open(\"tokenized_training_data.json\", 'w') as file:\n",
    "    json.dump(training_data, file, indent=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0d96a59",
   "metadata": {},
   "source": [
    "# Count how many qualified sentences there are"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f936f959-afbd-46e4-b798-08700974f41f",
   "metadata": {},
   "outputs": [],
   "source": [
    "max_token = 63\n",
    "min_token = 10\n",
    "counter = 0\n",
    "\n",
    "with open('tokenized_training_data.json', 'r') as file:\n",
    "    for item in ijson.items(file, 'item'):\n",
    "        for tokenized_sentence in item:\n",
    "            length = len(tokenized_sentence)\n",
    "            if length < max_token and length > min_token:\n",
    "                counter += 1\n",
    "\n",
    "print(counter)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75304ca0-0535-40cb-9f18-e5a784c73971",
   "metadata": {},
   "source": [
    "# Store tokenized data into tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b9975bf-17ac-47be-aaaa-0d5163e44c83",
   "metadata": {},
   "outputs": [],
   "source": [
    "shape = (117979663, 64)\n",
    "\n",
    "tokenized_tensor = torch.empty(shape, dtype=torch.int16)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9cb3b82-1fd4-494e-aab4-39c8151cf185",
   "metadata": {},
   "outputs": [],
   "source": [
    "max_token = 63\n",
    "min_token = 10\n",
    "counter = 0\n",
    "\n",
    "with open('tokenized_training_data.json', 'r') as file:\n",
    "    for item in ijson.items(file, 'item'):\n",
    "        for tokenized_sentence in item:\n",
    "            length = len(tokenized_sentence)\n",
    "            if length < max_token and length > min_token:\n",
    "                tokenized_sentence.append(tokenizer.eos_token_id)\n",
    "                tokenized_sentence = tokenized_sentence + [32000] * (64 - len(tokenized_sentence))\n",
    "                sentence_tokenized_tensor = torch.tensor(tokenized_sentence, dtype=torch.int16)\n",
    "                tokenized_tensor[counter] = sentence_tokenized_tensor\n",
    "                counter += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82721910-22cd-4adf-ad16-91a520d4627b",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(tokenized_tensor, \"llama2_wiki_64.pt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10640389",
   "metadata": {},
   "source": [
    "# Load token embedding and store it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8eddb729-1139-4d18-9eea-7a9f93a23889",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_id = \"NousResearch/Llama-2-7b-hf\"\n",
    "bnb_config = BitsAndBytesConfig(\n",
    "    load_in_4bit=True,\n",
    "    bnb_4bit_use_double_quant=True,\n",
    "    bnb_4bit_quant_type=\"nf4\",\n",
    "    bnb_4bit_compute_dtype=torch.bfloat16\n",
    ")\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=bnb_config)\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_id, padding_side=\"right\")\n",
    "tokenizer.pad_token_id = 128002\n",
    "\n",
    "# Access the embedding matrix\n",
    "word_embeddings_tensor = model.model.embed_tokens.weight.data\n",
    "\n",
    "# Delete llama3 because we are no longer using it.\n",
    "del model\n",
    "\n",
    "# Store vocabulary size and embedding dimension\n",
    "num_embeddings, embedding_dim = word_embeddings_tensor.shape\n",
    "word_embeddings_tensor.requires_grad = False\n",
    "\n",
    "torch.save(word_embeddings_tensor, 'word_embeddings_tensor_llama2.pt')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}