try_model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] ='0'\n",
    "\n",
    "from gpt2_model import GPT2, GPT2Tokenizer\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "dir_path = 'weights/fineweb-edu-15b/'\n",
    "tokenizer_path = dir_path + 'gpt2tokenizer_fineweb-edu.model'\n",
    "checkpoint_path = dir_path + 'model_27288.pt'\n",
    "checkpoint = torch.load(checkpoint_path, weights_only=False)\n",
    "\n",
    "device = 'cuda:0'\n",
    "\n",
    "tokenizer = GPT2Tokenizer.from_pretrained(tokenizer_path)\n",
    "model, _ = GPT2.from_checkpoint(checkpoint_path, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The capital of the United States is Washington D.C., and is now known as the capital of the United States. Washington D.C. is an urban center (which means it has \"urban\" in its language but not necessarily \"urban\" in any other way) with an average density of 413.8 people/ square feet, and is one of the most densely populated (and therefore the second most populous) metropolitan areas on the planet. With a population of 371,933,000, the United States is the most densely populated of the world's 48 contiguous states; its total population is 533,328. The metropolitan area of the United States is the largest in the world, covering an area of 4,690,511 square miles. For the year 2000, the Washington Square Park Resort and its other three attractions is the hub of the country's popular culture and is represented by more than 565 film films, television, radio and television stations, four radio networks,\n"
     ]
    }
   ],
   "source": [
    "prompt = 'The capital of the United States is'\n",
    "\n",
    "input_ids = tokenizer.encode(prompt, return_tensors=True, device=device)\n",
    "\n",
    "with torch.no_grad():\n",
    "    output = model.generate(input_ids, max_new_tokens=200, do_sample=True, topk=50, eot=tokenizer.special_tokens['<|endoftext|>'])\n",
    "\n",
    "print(tokenizer.decode(output))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "gpt2",
   "language": "python",
   "name": "gpt2"
  },
  "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}