01_Intro_LocalApp.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f9fdc8f2-9873-43ad-b3a0-cb26cfa54de1",
   "metadata": {},
   "source": [
    "# Run an LLM App in 15 Minutes\n",
    "---\n",
    "To prime ourselves for the type of work ahead, we will start by creating a [question answering (QA)](https://docs.langchain.com/docs/use-cases/qa-docs) service designed to run locally.\n",
    "\n",
    "Large language models (LLMs), while very impressive at next token prediction, have no relationship to the truth. This is especially relevant when the topic falls outside of the model's training data. To help mitigate their hallucinatory tendencies, we can implement a pattern referred to as [retrieval QA](https://python.langchain.com/en/latest/modules/chains/index_examples/vector_db_qa.html). In this use case, we generate embeddings for domain-specific documents that the LLM can then use to construct a response to a user query.\n",
    "\n",
    "After this short notebook, you will have set up a [document corpus](https://en.wikipedia.org/wiki/Text_corpus) of [Taylor Swift's Eras Tour](https://en.wikipedia.org/wiki/The_Eras_Tour) and the [2023 XFL Season](https://en.wikipedia.org/wiki/2023_XFL_season) for StableLM to use as context to supplement its generated answer."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "d37af2de",
   "metadata": {},
   "source": [
    "## Create a document corpus\n",
    "\n",
    "First, you need to establish the pool of information from which the language model will draw its context. In this example, we'll be using a few modules from [LangChain](https://python.langchain.com/en/latest/index.html) to facilitate this process.\n",
    "We'll be using [FAISS](https://engineering.fb.com/2017/03/29/data-infrastructure/faiss-a-library-for-efficient-similarity-search/), a library for similarity search across vector embeddings."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e0b67a33",
   "metadata": {},
   "source": [
    "### Load documents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "419dd205",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.document_loaders import WikipediaLoader\n",
    "\n",
    "topics = [\"The Eras Tour\", \"2023 XFL season\"]\n",
    "loaders = [WikipediaLoader(query=topic, load_max_docs=20) for topic in topics]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "a3b6fd85",
   "metadata": {},
   "source": [
    "### Split documents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fea8420f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "\n",
    "text_splitter = RecursiveCharacterTextSplitter(\n",
    "    chunk_size=300,\n",
    "    chunk_overlap=20,\n",
    "    length_function=len,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f1423a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from operator import add\n",
    "\n",
    "# Load documents\n",
    "docs = add(*[loader.load() for loader in loaders])\n",
    "print(\", \".join([d.metadata[\"title\"] for d in docs]))\n",
    "\n",
    "# Split documents into chunks\n",
    "chunks = text_splitter.create_documents(\n",
    "    [doc.page_content for doc in docs], metadatas=[doc.metadata for doc in docs]\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "8632b7db",
   "metadata": {},
   "source": [
    "### Create embeddings for documents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d3d83b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.embeddings.base import Embeddings\n",
    "from langchain.vectorstores import FAISS\n",
    "from sentence_transformers import SentenceTransformer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "797a228d-4722-4c0d-8b97-b8bb51a9f59b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class LocalHuggingFaceEmbeddings(Embeddings):\n",
    "    def __init__(self, model_id):\n",
    "        self.model = SentenceTransformer(model_id)\n",
    "\n",
    "    def embed_documents(self, texts: list[str]) -> list[list[float]]:\n",
    "        embeddings = self.model.encode(texts)\n",
    "        return embeddings\n",
    "\n",
    "    def embed_query(self, text: str) -> list[float]:\n",
    "        embedding = self.model.encode(text)\n",
    "        return list(map(float, embedding))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37d137cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings = LocalHuggingFaceEmbeddings(\"multi-qa-mpnet-base-dot-v1\")\n",
    "db = FAISS.from_documents(chunks, embeddings)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "55afa5e8",
   "metadata": {},
   "source": [
    "### Store documents and embeddings in a vector store"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "def90b66-216e-4f5b-b68d-c022c1b99a8f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "FAISS_INDEX_PATH = \"faiss_index_local\"\n",
    "\n",
    "db.save_local(FAISS_INDEX_PATH)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "30a4feef",
   "metadata": {},
   "source": [
    "## Set up a QA chain"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "7de66eef",
   "metadata": {},
   "source": [
    "### Create a custom `Pipeline`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "625306aa-195c-455b-b935-7081c86c33ab",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from langchain import HuggingFacePipeline\n",
    "from transformers import pipeline as hf_pipeline\n",
    "from typing import Optional, Any"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92bde69c-5642-4611-ab98-76cca879a137",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class StableLMPipeline(HuggingFacePipeline):\n",
    "    # Class is temporary, we are working with the authors of LangChain to make these unnecessary.\n",
    "\n",
    "    def _call(self, prompt: str, stop: Optional[list[str]] = None) -> str:\n",
    "        response = self.pipeline(\n",
    "            prompt, temperature=0.1, max_new_tokens=256, do_sample=True\n",
    "        )\n",
    "        print(f\"Response is: {response}\")\n",
    "        text = response[0][\"generated_text\"][len(prompt) :]\n",
    "        return text\n",
    "\n",
    "    @classmethod\n",
    "    def from_model_id(\n",
    "        cls,\n",
    "        model_id: str,\n",
    "        task: str,\n",
    "        device: Optional[str] = None,\n",
    "        model_kwargs: Optional[dict] = None,\n",
    "        **kwargs: Any,\n",
    "    ):\n",
    "        pipeline = hf_pipeline(\n",
    "            model=model_id,\n",
    "            task=task,\n",
    "            device=device,\n",
    "            model_kwargs=model_kwargs,\n",
    "        )\n",
    "        return cls(\n",
    "            pipeline=pipeline,\n",
    "            model_id=model_id,\n",
    "            model_kwargs=model_kwargs,\n",
    "            **kwargs,\n",
    "        )"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "2a3b1c6d",
   "metadata": {},
   "source": [
    "### Write a prompt template"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af642edf-f6e4-47d9-ad62-964298cc7414",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from langchain.prompts import PromptTemplate\n",
    "\n",
    "template = \"\"\"\n",
    "<|SYSTEM|># StableLM Tuned (Alpha version)\n",
    "- You are a helpful, polite, fact-based agent for answering questions. \n",
    "- Your answers include enough detail for someone to follow through on your suggestions. \n",
    "<|USER|>\n",
    "If you don't know the answer, just say that you don't know. Don't try to make up an answer.\n",
    "Please answer the following question using the context provided. \n",
    "\n",
    "CONTEXT: \n",
    "{context}\n",
    "=========\n",
    "QUESTION: {question} \n",
    "ANSWER: <|ASSISTANT|>\"\"\"\n",
    "\n",
    "PROMPT = PromptTemplate(template=template, input_variables=[\"context\", \"question\"])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "89dc7af5",
   "metadata": {},
   "source": [
    "### Create the QA chain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4b24b18",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from langchain.chains.question_answering import load_qa_chain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0dd2ba8-0d80-493e-8df4-b32f315a39a4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class QALocal:\n",
    "    def __init__(self):\n",
    "        self.embeddings = LocalHuggingFaceEmbeddings(\"multi-qa-mpnet-base-dot-v1\")\n",
    "        self.db = FAISS.load_local(FAISS_INDEX_PATH, self.embeddings)\n",
    "        self.llm = StableLMPipeline.from_model_id(\n",
    "            model_id=\"stabilityai/stablelm-tuned-alpha-7b\",\n",
    "            task=\"text-generation\",\n",
    "            model_kwargs={\n",
    "                \"torch_dtype\": torch.float16,\n",
    "                \"device_map\": \"auto\",\n",
    "                \"cache_dir\": \"/mnt/local_storage\",\n",
    "            },\n",
    "        )\n",
    "        self.chain = load_qa_chain(llm=self.llm, chain_type=\"stuff\", prompt=PROMPT)\n",
    "\n",
    "    def qa(self, query):\n",
    "        search_results = self.db.similarity_search(query)\n",
    "        print(f\"Results from db are: {search_results}\")\n",
    "        result = self.chain({\"input_documents\": search_results, \"question\": query})\n",
    "        print(f\"Result is: {result}\")\n",
    "        return result[\"output_text\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef1628de-7501-4a00-a0e0-7fe82b3d8619",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "local_qa = QALocal()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "bf5753e3",
   "metadata": {},
   "source": [
    "## Query the chain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98a49947-d2cd-40b3-8595-acfab33c6dcb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "local_qa.qa(\"How many people live in San Francisco?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0604df4f-1928-408c-a820-ee134ef72078",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "local_qa.qa(\"When did Taylor Swift's Eras tour start?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07e7d46e-d84f-4e68-9da7-344621805af7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "local_qa.qa(\"Can you tell me about the XFL 2023 season?\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "437c4b06",
   "metadata": {},
   "source": [
    "## Tear down application\n",
    "\n",
    "You can either shutdown the kernel or use these cells to free up memory occupied by this application."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0dc66fa8-5ea1-46dd-a6b7-cb4524af6d3a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "del local_qa"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "462df2c4-14c9-4a14-ae7a-8f10ccdbb51b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from accelerate import Accelerator\n",
    "\n",
    "accelerator = Accelerator()\n",
    "accelerator.free_memory()"
   ]
  }
 ],
 "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.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}