{
 "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 Chatbot\n",
    "from config import BASE_MODEL, MY_MODEL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "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": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Create chatbot instance using chat.py\n",
    "\"\"\"\n",
    "chatbot = Chatbot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Question: What options are available for someone in my situation?\n",
      "Response: I'm here to help you find suitable treatment options in the Boston area. Before we proceed, could you please share more information about your situation, such as:\n",
      "\n",
      "1. What type of substance(s) are you struggling with (e.g., opioids, alcohol, cocaine, etc.)?\n",
      "2. Do you have a mental health condition that you'd like to address in conjunction with substance use treatment (e.g., depression, anxiety, PTSD, etc.)?\n",
      "3. Are you looking for inpatient, outpatient, or intensive outpatient (IOP) treatment?\n",
      "4. Are you insured, and if so, what type of insurance do you have?\n",
      "5. Are you looking for a specific age range or demographic (e.g., young adults, older adults, family therapy, etc.)?\n",
      "6. Are there any specific treatment approaches or modalities you're interested in (e.g., medication-assisted treatment, cognitive-behavioral therapy, etc.)?\n",
      "\n",
      "The more information you provide, the better I can tailor my recommendations to suit your needs. \n",
      "\n",
      "Here are some Boston area treatment facilities that may be a good starting point:\n",
      "\n",
      "1. Boston Medical Center's (BMC) Gosnold: Offers comprehensive substance use and mental health services, including inpatient and outpatient programs.\n",
      "2. Caritas HealthCare: Provides a range of services, including medication-assisted treatment, counseling, and case management.\n",
      "3. McLean Hospital: A leading mental health facility that offers a range of inpatient and outpatient services, including substance use treatment.\n",
      "4. The Center for Psychiatric Rehabilitation (CPR): Offers a comprehensive range of services, including substance use treatment, counseling, and vocational support.\n",
      "5. Horizon House: Provides a range of services, including inpatient and outpatient substance use treatment, as well as mental health services.\n",
      "6. New England Baptist Hospital's (NEBH) Behavioral Health Services: Offers a range of services, including inpatient and outpatient substance use treatment.\n",
      "\n",
      "Please let me know if you have any specific questions or concerns, or if there's anything else I can do to assist you in finding the right treatment facility for your needs.\n"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "Test out generating some responses from the chatbot.\n",
    "Inference time\n",
    "\"\"\"\n",
    "test_question = \"What options are available for someone in my situation?\"\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 task-specific 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 relevant 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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}