{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Usage\n",
    "1. Install python dependencies\n",
    "```shell\n",
    "!pip install pypdf langchain unstructured transformers_stream_generator\n",
    "!pip install modelscope  nltk pydantic  tiktoken  llama-index\n",
    "```\n",
    "\n",
    "2. Download data files we need in this example\n",
    "```shell\n",
    "!wget https://modelscope.oss-cn-beijing.aliyuncs.com/resource/rag/averaged_perceptron_tagger.zip\n",
    "!wget https://modelscope.oss-cn-beijing.aliyuncs.com/resource/rag/punkt.zip\n",
    "!wget https://modelscope.oss-cn-beijing.aliyuncs.com/resource/rag/xianjiaoda.md\n",
    "\n",
    "!mkdir -p /root/nltk_data/tokenizers\n",
    "!mkdir -p /root/nltk_data/taggers\n",
    "!cp /mnt/workspace/punkt.zip /root/nltk_data/tokenizers\n",
    "!cp /mnt/workspace/averaged_perceptron_tagger.zip /root/nltk_data/taggers\n",
    "!cd /root/nltk_data/tokenizers; unzip punkt.zip;\n",
    "!cd /root/nltk_data/taggers; unzip averaged_perceptron_tagger.zip;\n",
    "\n",
    "!mkdir -p /mnt/workspace/custom_data\n",
    "!mv /mnt/workspace/xianjiaoda.md /mnt/workspace/custom_data\n",
    "\n",
    "!cd /mnt/workspace\n",
    "``` \n",
    "\n",
    "3. Enjoy your QA AI"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "8230365523c9330a"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a407764-9392-48ae-9bed-8c73c9f76fbc",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-16T08:58:56.323000Z",
     "iopub.status.busy": "2024-01-16T08:58:56.322690Z",
     "iopub.status.idle": "2024-01-16T08:59:57.862755Z",
     "shell.execute_reply": "2024-01-16T08:59:57.862041Z",
     "shell.execute_reply.started": "2024-01-16T08:58:56.322980Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "!pip install pypdf langchain unstructured transformers_stream_generator\n",
    "!pip install modelscope  nltk pydantic  tiktoken  llama-index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "696c6b78-53e8-4135-8376-ce8902b7d79a",
   "metadata": {
    "ExecutionIndicator": {
     "show": true
    },
    "execution": {
     "iopub.execute_input": "2024-01-16T09:04:59.193375Z",
     "iopub.status.busy": "2024-01-16T09:04:59.193082Z",
     "iopub.status.idle": "2024-01-16T09:05:00.971449Z",
     "shell.execute_reply": "2024-01-16T09:05:00.970857Z",
     "shell.execute_reply.started": "2024-01-16T09:04:59.193357Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "!wget https://modelscope.oss-cn-beijing.aliyuncs.com/resource/rag/averaged_perceptron_tagger.zip\n",
    "!wget https://modelscope.oss-cn-beijing.aliyuncs.com/resource/rag/punkt.zip\n",
    "!wget https://modelscope.oss-cn-beijing.aliyuncs.com/resource/rag/xianjiaoda.md\n",
    "\n",
    "!mkdir -p /root/nltk_data/tokenizers\n",
    "!mkdir -p /root/nltk_data/taggers\n",
    "!cp /mnt/workspace/punkt.zip /root/nltk_data/tokenizers\n",
    "!cp /mnt/workspace/averaged_perceptron_tagger.zip /root/nltk_data/taggers\n",
    "!cd /root/nltk_data/tokenizers; unzip punkt.zip;\n",
    "!cd /root/nltk_data/taggers; unzip averaged_perceptron_tagger.zip;\n",
    "\n",
    "!mkdir -p /mnt/workspace/custom_data\n",
    "!mv /mnt/workspace/xianjiaoda.md /mnt/workspace/custom_data\n",
    "\n",
    "!cd /mnt/workspace"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1cb8feca-c71f-4ad6-8eff-caae95411aa0",
   "metadata": {
    "ExecutionIndicator": {
     "show": true
    },
    "execution": {
     "iopub.execute_input": "2024-01-16T09:06:03.024995Z",
     "iopub.status.busy": "2024-01-16T09:06:03.024622Z",
     "iopub.status.idle": "2024-01-16T09:09:15.894774Z",
     "shell.execute_reply": "2024-01-16T09:09:15.894230Z",
     "shell.execute_reply.started": "2024-01-16T09:06:03.024974Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from abc import ABC\n",
    "from typing import Any, List, Optional, Dict, cast\n",
    "\n",
    "import torch\n",
    "from langchain_core.language_models.llms import LLM\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_core.runnables import RunnablePassthrough\n",
    "from modelscope import AutoModelForCausalLM, AutoTokenizer\n",
    "from llama_index import GPTVectorStoreIndex, SimpleDirectoryReader\n",
    "from llama_index import ServiceContext\n",
    "from llama_index.embeddings.base import BaseEmbedding\n",
    "from llama_index import set_global_service_context\n",
    "from langchain_core.retrievers import BaseRetriever\n",
    "from langchain_core.callbacks import CallbackManagerForRetrieverRun\n",
    "from langchain_core.documents import Document\n",
    "from llama_index.retrievers import VectorIndexRetriever\n",
    "\n",
    "# configs for LLM\n",
    "llm_name = \"Qwen/Qwen-1_8B-Chat\"\n",
    "llm_revision = \"master\"\n",
    "\n",
    "# configs for embedding model\n",
    "embedding_model = \"damo/nlp_gte_sentence-embedding_chinese-small\"\n",
    "\n",
    "# file path for your custom knowledge base\n",
    "knowledge_doc_file_dir = \"/mnt/workspace/custom_data/\"\n",
    "knowledge_doc_file_path = knowledge_doc_file_dir + \"xianjiaoda.md\"\n",
    "\n",
    "\n",
    "# define our Embedding class to use models in Modelscope\n",
    "class ModelScopeEmbeddings4LlamaIndex(BaseEmbedding, ABC):\n",
    "    embed: Any = None\n",
    "    model_id: str = \"damo/nlp_gte_sentence-embedding_chinese-small\"\n",
    "\n",
    "    def __init__(\n",
    "            self,\n",
    "            model_id: str,\n",
    "            **kwargs: Any,\n",
    "    ) -> None:\n",
    "        super().__init__(**kwargs)\n",
    "        try:\n",
    "            from modelscope.models import Model\n",
    "            from modelscope.pipelines import pipeline\n",
    "            from modelscope.utils.constant import Tasks\n",
    "            self.embed = pipeline(Tasks.sentence_embedding, model=self.model_id)\n",
    "\n",
    "        except ImportError as e:\n",
    "            raise ValueError(\n",
    "                \"Could not import some python packages.\" \"Please install it with `pip install modelscope`.\"\n",
    "            ) from e\n",
    "\n",
    "    def _get_query_embedding(self, query: str) -> List[float]:\n",
    "        text = query.replace(\"\\n\", \" \")\n",
    "        inputs = {\"source_sentence\": [text]}\n",
    "        return self.embed(input=inputs)['text_embedding'][0]\n",
    "\n",
    "    def _get_text_embedding(self, text: str) -> List[float]:\n",
    "        text = text.replace(\"\\n\", \" \")\n",
    "        inputs = {\"source_sentence\": [text]}\n",
    "        return self.embed(input=inputs)['text_embedding'][0]\n",
    "\n",
    "    def _get_text_embeddings(self, texts: List[str]) -> List[List[float]]:\n",
    "        texts = list(map(lambda x: x.replace(\"\\n\", \" \"), texts))\n",
    "        inputs = {\"source_sentence\": texts}\n",
    "        return self.embed(input=inputs)['text_embedding']\n",
    "\n",
    "    async def _aget_query_embedding(self, query: str) -> List[float]:\n",
    "        return self._get_query_embedding(query)\n",
    "\n",
    "\n",
    "# define our Retriever with llama-index to co-operate with Langchain\n",
    "# note that the 'LlamaIndexRetriever' defined in langchain-community.retrievers.llama_index.py\n",
    "# is no longer compatible with llamaIndex code right now.\n",
    "class LlamaIndexRetriever(BaseRetriever):\n",
    "    index: Any\n",
    "    \"\"\"LlamaIndex index to query.\"\"\"\n",
    "\n",
    "    def _get_relevant_documents(\n",
    "        self, query: str, *, run_manager: CallbackManagerForRetrieverRun\n",
    "    ) -> List[Document]:\n",
    "        \"\"\"Get documents relevant for a query.\"\"\"\n",
    "        try:\n",
    "            from llama_index.indices.base import BaseIndex\n",
    "            from llama_index.response.schema import Response\n",
    "        except ImportError:\n",
    "            raise ImportError(\n",
    "                \"You need to install `pip install llama-index` to use this retriever.\"\n",
    "            )\n",
    "        index = cast(BaseIndex, self.index)\n",
    "        print('@@@ query=', query)\n",
    "\n",
    "        response = index.as_query_engine().query(query)\n",
    "        response = cast(Response, response)\n",
    "        # parse source nodes\n",
    "        docs = []\n",
    "        for source_node in response.source_nodes:\n",
    "            print('@@@@ source=', source_node)\n",
    "            metadata = source_node.metadata or {}\n",
    "            docs.append(\n",
    "                Document(page_content=source_node.get_text(), metadata=metadata)\n",
    "            )\n",
    "        return docs\n",
    "\n",
    "def torch_gc():\n",
    "    os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n",
    "    DEVICE = \"cuda\"\n",
    "    DEVICE_ID = \"0\"\n",
    "    CUDA_DEVICE = f\"{DEVICE}:{DEVICE_ID}\" if DEVICE_ID else DEVICE\n",
    "    a = torch.Tensor([1, 2])\n",
    "    a = a.cuda()\n",
    "    print(a)\n",
    "\n",
    "    if torch.cuda.is_available():\n",
    "        with torch.cuda.device(CUDA_DEVICE):\n",
    "            torch.cuda.empty_cache()\n",
    "            torch.cuda.ipc_collect()\n",
    "\n",
    "\n",
    "# global resources used by QianWenChatLLM (this is not a good practice)\n",
    "tokenizer = AutoTokenizer.from_pretrained(llm_name, revision=llm_revision, trust_remote_code=True)\n",
    "model = AutoModelForCausalLM.from_pretrained(llm_name, revision=llm_revision, device_map=\"auto\",\n",
    "                                             trust_remote_code=True, fp16=True).eval()\n",
    "\n",
    "\n",
    "# define QianWen LLM based on langchain's LLM to use models in Modelscope\n",
    "class QianWenChatLLM(LLM):\n",
    "    max_length = 10000\n",
    "    temperature: float = 0.01\n",
    "    top_p = 0.9\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "    @property\n",
    "    def _llm_type(self):\n",
    "        return \"ChatLLM\"\n",
    "\n",
    "    def _call(\n",
    "            self,\n",
    "            prompt: str,\n",
    "            stop: Optional[List[str]] = None,\n",
    "            run_manager=None,\n",
    "            **kwargs: Any,\n",
    "    ) -> str:\n",
    "        print(prompt)\n",
    "        response, history = model.chat(tokenizer, prompt, history=None)\n",
    "        torch_gc()\n",
    "        return response\n",
    "\n",
    "\n",
    "# STEP1: create LLM instance\n",
    "qwllm = QianWenChatLLM()\n",
    "print('STEP1: qianwen LLM created')\n",
    "\n",
    "# STEP2: load knowledge file and initialize vector db by llamaIndex\n",
    "print('STEP2: reading docs ...')\n",
    "embeddings = ModelScopeEmbeddings4LlamaIndex(model_id=embedding_model)\n",
    "service_context = ServiceContext.from_defaults(embed_model=embeddings, llm=None)\n",
    "set_global_service_context(service_context)     # global config, not good\n",
    "\n",
    "llamaIndex_docs = SimpleDirectoryReader(knowledge_doc_file_dir).load_data()\n",
    "llamaIndex_index = GPTVectorStoreIndex.from_documents(llamaIndex_docs, chunk_size=512)\n",
    "retriever = LlamaIndexRetriever(index=llamaIndex_index)\n",
    "print(' 2.2 reading doc done, vec db created.')\n",
    "\n",
    "# STEP3: create chat template\n",
    "prompt_template = \"\"\"请基于```内的内容回答问题。\"\n",
    "```\n",
    "{context}\n",
    "```\n",
    "我的问题是：{question}。\n",
    "\"\"\"\n",
    "prompt = ChatPromptTemplate.from_template(template=prompt_template)\n",
    "print('STEP3: chat prompt template created.')\n",
    "\n",
    "# STEP4: create RAG chain to do QA\n",
    "chain = (\n",
    "        {\"context\": retriever, \"question\": RunnablePassthrough()}\n",
    "        | prompt\n",
    "        | qwllm\n",
    "        | StrOutputParser()\n",
    ")\n",
    "chain.invoke('西安交大的校训是什么？')\n",
    "# chain.invoke('魔搭社区有哪些模型?')\n",
    "# chain.invoke('modelscope是什么?')\n",
    "# chain.invoke('萧峰和乔峰是什么关系?')\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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
