final.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_cell_guid": "b1076dfc-b9ad-4769-8c92-a6c4dae69d19",
    "_uuid": "8f2839f25d086af736a60e9eeb907d3b93b6e0e5"
   },
   "outputs": [],
   "source": [
    "!pip install rotary_embedding_torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import wandb\n",
    "\n",
    "wandb.login(key='')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from transformers import GPT2LMHeadModel, GPT2Tokenizer, TrainingArguments, Trainer, GPT2Config\n",
    "from datasets import load_dataset\n",
    "from rotary_embedding_torch import RotaryEmbedding\n",
    "from huggingface_hub import HfApi\n",
    "\n",
    "max_length = 4096\n",
    "\n",
    "model_name = \"gpt2-medium\"\n",
    "config = GPT2Config.from_pretrained(model_name)\n",
    "config.max_position_embeddings = max_length\n",
    "config.n_ctx = max_length\n",
    "\n",
    "model = GPT2LMHeadModel(config)\n",
    "tokenizer = GPT2Tokenizer.from_pretrained(model_name)\n",
    "if tokenizer.pad_token is None:\n",
    "    tokenizer.add_special_tokens({'pad_token': tokenizer.eos_token})\n",
    "\n",
    "rotary_emb = RotaryEmbedding(\n",
    "    dim=32,\n",
    "    interpolate_factor=2.0,\n",
    ")\n",
    "\n",
    "for layer in model.transformer.h:\n",
    "    layer.attn.rotary_emb = rotary_emb\n",
    "\n",
    "for name, param in model.named_parameters():\n",
    "    if \"rotary_emb\" not in name and \"ln_f\" not in name:\n",
    "        param.requires_grad = False\n",
    "\n",
    "dataset = load_dataset(\"Yukang/LongAlpaca-12k\", split='train[:10%]')\n",
    "\n",
    "def tokenize_function(examples):\n",
    "    tokenized = tokenizer(examples['output'], truncation=True, max_length=max_length, padding=\"max_length\", return_tensors=\"pt\")\n",
    "    tokenized[\"labels\"] = tokenized[\"input_ids\"].clone()\n",
    "    return tokenized\n",
    "\n",
    "tokenized_datasets = dataset.map(tokenize_function, batched=True, remove_columns=[\"input\", \"output\"])\n",
    "\n",
    "class CustomTrainer(Trainer):\n",
    "    def compute_loss(self, model, inputs, return_outputs=False):\n",
    "        labels = inputs.pop(\"labels\")\n",
    "        outputs = model(**inputs, labels=labels)\n",
    "        return (outputs.loss, outputs) if return_outputs else outputs.loss\n",
    "\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=\"./results\",\n",
    "    evaluation_strategy=\"steps\",\n",
    "    eval_steps=10,\n",
    "    learning_rate=2e-5,\n",
    "    per_device_train_batch_size=1,\n",
    "    per_device_eval_batch_size=1,\n",
    "    num_train_epochs=1,\n",
    "    weight_decay=0.01,\n",
    "    max_steps=1000,\n",
    "    logging_dir='./logs',\n",
    "    logging_steps=10,\n",
    ")\n",
    "\n",
    "trainer = CustomTrainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=tokenized_datasets,\n",
    "    eval_dataset=tokenized_datasets,\n",
    ")\n",
    "\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import login\n",
    "\n",
    "login(token=\"\")\n",
    "repo_name = \"archit11/final_gpt2\"  \n",
    "\n",
    "model.push_to_hub(repo_name)\n",
    "tokenizer.push_to_hub(repo_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import GPT2LMHeadModel, GPT2Tokenizer\n",
    "\n",
    "model_path = \"./results\"  \n",
    "tokenizer = GPT2Tokenizer.from_pretrained(model_path)\n",
    "model = GPT2LMHeadModel.from_pretrained(model_path)\n",
    "\n",
    "# Set the model to evaluation mode\n",
    "model.eval()\n",
    "\n",
    "def generate_text(prompt, max_length=100):\n",
    "    input_ids = tokenizer.encode(prompt, return_tensors=\"pt\")\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        output = model.generate(\n",
    "            input_ids,\n",
    "            max_length=max_length,\n",
    "            num_return_sequences=1,\n",
    "            no_repeat_ngram_size=2,\n",
    "            top_k=50,\n",
    "            top_p=0.95,\n",
    "            temperature=0.7\n",
    "        )\n",
    "    \n",
    "    return tokenizer.decode(output[0], skip_special_tokens=True)\n",
    "\n",
    "prompt = \"Once upon a time\"\n",
    "generated_text = generate_text(prompt)\n",
    "print(f\"Generated text:\\n{generated_text}\")\n",
    "\n",
    "save_directory = \"./saved_model\"\n",
    "model.save_pretrained(save_directory)\n",
    "tokenizer.save_pretrained(save_directory)\n",
    "\n",
    "print(f\"Model and tokenizer saved to {save_directory}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "from huggingface_hub import login\n",
    "\n",
    "login(token=\"\")\n",
    "repo_name = \"archit11/final_gpt2\"  # Replace with your desired model name\n",
    "\n",
    "# Push the model and tokenizer to the Hub\n",
    "model.push_to_hub(repo_name)\n",
    "tokenizer.push_to_hub(repo_name)"
   ]
  }
 ],
 "metadata": {
  "kaggle": {
   "accelerator": "nvidiaTeslaT4",
   "dataSources": [],
   "isGpuEnabled": true,
   "isInternetEnabled": true,
   "language": "python",
   "sourceType": "notebook"
  },
  "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}