{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chatbot Development\n",
    "\n",
    "Use this notebook to load the model and then initialize, update, and test the chatbot."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setup and Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from huggingface_hub import login\n",
    "\n",
    "\n",
    "from src.chat import SchoolChatbot\n",
    "from config import BASE_MODEL, MY_MODEL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "50617a1a9bdc434d9b897a5e9b529b52",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(HTML(value='<center> <img\\nsrc=https://huggingface.co/front/assets/huggingface_logo-noborder.sv…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\"\"\"\n",
    "TODO: Add your Hugging Face token\n",
    "Options:\n",
    "1. Use login() and enter token when prompted. It won't ask for your token if you already logged in using the command: huggingface-cli login in the terminal.\n",
    "2. Set environment variable HUGGINGFACE_TOKEN\n",
    "3. Pass token directly (not recommended for shared notebooks)\n",
    "\"\"\"\n",
    "\n",
    "login()\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initialize and test chatbot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Create chatbot instance using chat.py\n",
    "\"\"\"\n",
    "chatbot = SchoolChatbot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Question: I live in Jamaica Plain and want to send my child to a school that offers Spanish classes. What schools are available?\n",
      "Response: There are several Boston Public Schools in and around Jamaica Plain that offer Spanish language instruction. Here are some options:\n",
      "\n",
      "1. Joseph P. Kelly Elementary School: This school offers Spanish bilingual education for students in grades Pre-K through 5. It is located at 1200 Columbus Ave, Boston, MA 02130.\n",
      "\n",
      "2. Edwin M. Bacon Elementary School: This school offers Spanish language instruction in grades Pre-K through 6. It is located at 30 Fayette St, Boston, MA 02130.\n",
      "\n",
      "3. Samuel J. & Dorothy Donnelly Elementary School: This school offers a Spanish Dual Language Learning Program in grades Pre-K through 5. It is located at 524 Cummins Highway, Roxbury, MA 02119.\n",
      "\n",
      "4. James Otis Elementary School: This school offers a Spanish Bilingual Program in grades Pre-K through 6. It is located at 1234 Columbus Ave, Boston, MA 02120.\n",
      "\n",
      "5. The John D. O'Bryant School of Mathematics and Science: This middle school (grades 6-8) offers Spanish language instruction as part of its academic program. It is located at 1345 Tremont St, Roxbury, MA 02120.\n",
      "\n",
      "It's recommended that you contact these schools directly to learn more about their programs and enrollment processes. Good luck!\n"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "Test out generating some responses from the chatbot.\n",
    "Inference time\n",
    "\"\"\"\n",
    "test_question = \"I live in Jamaica Plain and want to send my child to a school that offers Spanish classes. What schools are available?\"\n",
    "\n",
    "print(f\"\\nQuestion: {test_question}\")\n",
    "response = chatbot.get_response(test_question)\n",
    "print(f\"Response: {response}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TODO: Update pre-trained Llama to be a school choice chatbot\n",
    "\n",
    "This part is up to you! You might want to finetune the model, simply make a really good system prompt, use RAG, provide the model boston school choice data in-context, etc. Be creative!\n",
    "\n",
    "You can also feel free to do this in another script and then evaluate the model here.\n",
    "\n",
    "Tips:\n",
    "- HuggingFace has built-in methods to finetune models, if you choose that route. Take advantage of those methods! You can then save your new, finetuned model in the HuggingFace Hub. Change MY_MODEL in config.py to the name of the model in the hub to make your chatbot use it.\n",
    "- You may also want to consider LoRA if you choose finetuning."
   ]
  }
 ],
 "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}