{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Memorize\n",
    "\n",
    "Fine-tuning LLM itself to memorize information using unsupervised learning.\n",
    "\n",
    "This tool requires LLMs that support fine-tuning. Currently, only `langchain.llms import GradientLLM` is supported."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "from langchain.agents import AgentExecutor, AgentType, initialize_agent, load_tools\n",
    "from langchain.chains import LLMChain\n",
    "from langchain.memory import ConversationBufferMemory\n",
    "from langchain_community.llms import GradientLLM"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set the Environment API Key\n",
    "Make sure to get your API key from Gradient AI. You are given $10 in free credits to test and fine-tune different models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "from getpass import getpass\n",
    "\n",
    "if not os.environ.get(\"GRADIENT_ACCESS_TOKEN\", None):\n",
    "    # Access token under https://auth.gradient.ai/select-workspace\n",
    "    os.environ[\"GRADIENT_ACCESS_TOKEN\"] = getpass(\"gradient.ai access token:\")\n",
    "if not os.environ.get(\"GRADIENT_WORKSPACE_ID\", None):\n",
    "    # `ID` listed in `$ gradient workspace list`\n",
    "    # also displayed after login at at https://auth.gradient.ai/select-workspace\n",
    "    os.environ[\"GRADIENT_WORKSPACE_ID\"] = getpass(\"gradient.ai workspace id:\")\n",
    "if not os.environ.get(\"GRADIENT_MODEL_ADAPTER_ID\", None):\n",
    "    # `ID` listed in `$ gradient model list --workspace-id \"$GRADIENT_WORKSPACE_ID\"`\n",
    "    os.environ[\"GRADIENT_MODEL_ID\"] = getpass(\"gradient.ai model id:\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Optional: Validate your Environment variables ```GRADIENT_ACCESS_TOKEN``` and ```GRADIENT_WORKSPACE_ID``` to get currently deployed models."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create the `GradientLLM` instance\n",
    "You can specify different parameters such as the model name, max tokens generated, temperature, etc."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = GradientLLM(\n",
    "    model_id=os.environ[\"GRADIENT_MODEL_ID\"],\n",
    "    # # optional: set new credentials, they default to environment variables\n",
    "    # gradient_workspace_id=os.environ[\"GRADIENT_WORKSPACE_ID\"],\n",
    "    # gradient_access_token=os.environ[\"GRADIENT_ACCESS_TOKEN\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "tools = load_tools([\"memorize\"], llm=llm)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initiate the Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "agent = initialize_agent(\n",
    "    tools,\n",
    "    llm,\n",
    "    agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,\n",
    "    verbose=True,\n",
    "    # memory=ConversationBufferMemory(memory_key=\"chat_history\", return_messages=True),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run the agent\n",
    "Ask the agent to memorize a piece of text."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new AgentExecutor chain...\u001b[0m\n",
      "\u001b[32;1m\u001b[1;3mI should memorize this fact.\n",
      "Action: Memorize\n",
      "Action Input: Zara T\u001b[0m\n",
      "Observation: \u001b[36;1m\u001b[1;3mTrain complete. Loss: 1.6853971333333335\u001b[0m\n",
      "Thought:\u001b[32;1m\u001b[1;3mI now know the final answer.\n",
      "Final Answer: Zara Tubikova set a world\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'Zara Tubikova set a world'"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent.run(\n",
    "    \"Please remember the fact in detail:\\nWith astonishing dexterity, Zara Tubikova set a world record by solving a 4x4 Rubik's Cube variation blindfolded in under 20 seconds, employing only their feet.\"\n",
    ")"
   ]
  }
 ],
 "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.11.6"
  },
  "vscode": {
   "interpreter": {
    "hash": "a0a0263b650d907a3bfe41c0f8d6a63a071b884df3cfdc1579f00cdc1aed6b03"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}