{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "torch.cuda.is_available()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check that we can load the tokenizer and the model.  The first time this runs it will take a while.  The files go under ~/.cache/huggingface"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The tokenizer class you load from this checkpoint is not the same type as the class this function is called from. It may result in unexpected tokenization. \n",
      "The tokenizer class you load from this checkpoint is 'LLaMATokenizer'. \n",
      "The class this function is called from is 'LlamaTokenizer'.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fc0f4312aa7c4009a912c66dd1443763",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/33 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from transformers import LlamaTokenizer, LlamaForCausalLM\n",
    "tokenizer = LlamaTokenizer.from_pretrained(\n",
    "    \"decapoda-research/llama-7b-hf\")\n",
    "   \n",
    "#tokenizer.pad_token_id = (0)\n",
    "#tokenizer.padding_side = 'left'\n",
    "model = LlamaForCausalLM.from_pretrained(\n",
    "    \"decapoda-research/llama-7b-hf\",\n",
    "    device_map=\"auto\",\n",
    "    torch_dtype=torch.float16\n",
    ")\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Construct a basic example to interact with the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/matt/hf/sqllama-V0/.venv/lib/python3.7/site-packages/transformers/generation/utils.py:1220: UserWarning: You have modified the pretrained model configuration to control generation. This is a deprecated strategy to control generation and will be removed soon, in a future version. Please use a generation configuration file (see https://huggingface.co/docs/transformers/main_classes/text_generation)\n",
      "  \"You have modified the pretrained model configuration to control generation. This is a\"\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " ⁇  hey dude, talk to me.\n",
      "I'm a 20 year old guy from the UK. I'm a bit of a nerd, I like to read, I like to write, I like to play video games, I like to watch movies, I like to listen\n",
      " ⁇  whats the capital of georgia?\n",
      "What is the capital of Georgia?\n",
      "The capital of Georgia is Atlanta.\n",
      "What is the capital of Georgia?\n",
      "The capital of Georgia is Atlanta. The capital of Georgia is Atlanta. The capital of Georgia is Atlanta. The capital of Georgia is Atlanta. The\n"
     ]
    }
   ],
   "source": [
    "def ask(q,l=64):\n",
    "    toks = tokenizer(q , return_tensors='pt')\n",
    "    ctoks = toks.input_ids.to('cuda')\n",
    "    gen = model.generate(ctoks, max_length=64)\n",
    "    return tokenizer.decode(gen[0])\n",
    "\n",
    "r = ask('hey dude, talk to me')\n",
    "print(r)\n",
    "\n",
    "r = ask('whats the capital of georgia?')\n",
    "print(r)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get the WikiSQL project so we can get the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fatal: destination path 'WikiSQL' already exists and is not an empty directory.\n",
      "data/\n",
      "data/train.jsonl\n",
      "data/test.tables.jsonl\n",
      "data/test.db\n",
      "data/dev.tables.jsonl\n",
      "data/dev.db\n",
      "data/test.jsonl\n",
      "data/train.tables.jsonl\n",
      "data/train.db\n",
      "data/dev.jsonl\n"
     ]
    }
   ],
   "source": [
    "! git clone https://github.com/salesforce/WikiSQL\n",
    "! tar xvjf WikiSQL/data.tar.bz2"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.7.3"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "6a381460736e8a0eabfb35eafae436ba15c06439de44e28b965ea473bd8dda90"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}