{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Remote Execution"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In its current state, `NDIF` requires you to receive an API key. To get one, simply\n",
    "go to https://login.ndif.us and sign up.\n",
    "\n",
    "With a valid API key, you then can configure `nnsight` by doing the following:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nnsight import CONFIG\n",
    "\n",
    "CONFIG.set_default_api_key(\"YOUR_API_KEY\") \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This only needs to be run once as it will save this api key as the default in a\n",
    "config file along with the `nnsight` installation.\n",
    "\n",
    "Let's demonstrate using `nnsight`'s tracing\n",
    "context with one of the larger open source language models, `Llama-3.1-70b`!\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# llama3.1 70b is a gated model and you need access via your huggingface token\n",
    "os.environ['HF_TOKEN'] = \"YOUR_HUGGING_FACE_TOKEN\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-08-30 07:11:21,150 MainProcess nnsight_remote INFO     36ff46f0-d81a-4586-b7e7-eaf6f97d6c0b - RECEIVED: Your job has been received and is waiting approval.\n",
      "2024-08-30 07:11:21,184 MainProcess nnsight_remote INFO     36ff46f0-d81a-4586-b7e7-eaf6f97d6c0b - APPROVED: Your job was approved and is waiting to be run.\n",
      "2024-08-30 07:11:21,206 MainProcess nnsight_remote INFO     36ff46f0-d81a-4586-b7e7-eaf6f97d6c0b - RUNNING: Your job has started running.\n",
      "2024-08-30 07:11:21,398 MainProcess nnsight_remote INFO     36ff46f0-d81a-4586-b7e7-eaf6f97d6c0b - COMPLETED: Your job has been completed.\n",
      "Downloading result:   0%|          | 0.00/9.48M [00:00<?, ?B/s]huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n",
      "Downloading result: 100%|██████████| 9.48M/9.48M [00:02<00:00, 3.21MB/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(tensor([[[ 5.4688, -4.9062,  2.2344,  ..., -3.6875,  0.9609,  1.2578],\n",
      "         [ 1.5469, -0.6172, -1.4531,  ..., -1.1562, -0.1406, -2.1250],\n",
      "         [ 1.7812, -1.8906, -1.1875,  ...,  0.1680,  0.9609,  0.5625],\n",
      "         ...,\n",
      "         [ 0.9453, -0.3711,  1.3516,  ...,  1.3828, -0.7969, -1.9297],\n",
      "         [-0.8906,  0.3672,  0.2617,  ...,  2.4688, -0.4414, -0.6758],\n",
      "         [-1.6094,  1.0938,  1.7031,  ...,  1.8672, -1.1328, -0.5000]]],\n",
      "       dtype=torch.bfloat16), DynamicCache())\n",
      "tensor([[[ 6.3750,  8.6250, 13.0000,  ..., -4.1562, -4.1562, -4.1562],\n",
      "         [-2.8594, -2.2344, -3.0938,  ..., -8.6250, -8.6250, -8.6250],\n",
      "         [ 8.9375,  3.5938,  4.5000,  ..., -3.9375, -3.9375, -3.9375],\n",
      "         ...,\n",
      "         [ 3.5781,  3.4531,  0.0796,  ..., -6.5625, -6.5625, -6.5625],\n",
      "         [10.8750,  6.4062,  4.9375,  ..., -4.0000, -4.0000, -3.9844],\n",
      "         [ 7.2500,  6.1562,  3.5156,  ..., -4.7188, -4.7188, -4.7188]]])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# We'll never actually load the parameters so no need to specify a device_map.\n",
    "llama = LanguageModel(\"meta-llama/Meta-Llama-3.1-70B\")\n",
    "\n",
    "# All we need to specify using NDIF vs executing locally is remote=True.\n",
    "with llama.trace(\"The Eiffel Tower is in the city of\", remote=True) as runner:\n",
    "\n",
    "    hidden_states = llama.model.layers[-1].output.save()\n",
    "\n",
    "    output = llama.output.save()\n",
    "\n",
    "print(hidden_states)\n",
    "\n",
    "print(output[\"logits\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It really is as simple as `remote=True`. All of the techniques available in NNsight locally work just the same when running remotely."
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
