{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Usage\n",
    "\n",
    "## 1. Install necessary libs\n",
    "```shell\n",
    "!pip install modelscope\n",
    "!pip install transformers -U\n",
    "!pip install llama-index llama-index-llms-huggingface ipywidgets \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/punkt.zip\n",
    "!wget https://modelscope.oss-cn-beijing.aliyuncs.com/resource/rag/stopwords.zip\n",
    "!wget https://modelscope.oss-cn-beijing.aliyuncs.com/resource/rag/xianjiaoda.md\n",
    "\n",
    "!mkdir -p /opt/conda/lib/python3.10/site-packages/llama_index/core/_static/nltk_cache/tokenizers\n",
    "!mkdir -p /opt/conda/lib/python3.10/site-packages/llama_index/core/_static/nltk_cache/corpora\n",
    "\n",
    "!cp /mnt/workspace/punkt.zip /opt/conda/lib/python3.10/site-packages/llama_index/core/_static/nltk_cache/tokenizers\n",
    "!cp /mnt/workspace/stopwords.zip /opt/conda/lib/python3.10/site-packages/llama_index/core/_static/nltk_cache/corpora\n",
    "!cd /opt/conda/lib/python3.10/site-packages/llama_index/core/_static/nltk_cache/tokenizers; unzip punkt.zip;\n",
    "!cd /opt/conda/lib/python3.10/site-packages/llama_index/core/_static/nltk_cache/corpora; unzip stopwords.zip;\n",
    "\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. Go!"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "f4abc589d9bfffca"
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "!pip install modelscope\n",
    "!pip install transformers -U\n",
    "!pip install llama-index llama-index-llms-huggingface ipywidgets "
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "c32122833dd7b8c8"
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "!wget https://modelscope.oss-cn-beijing.aliyuncs.com/resource/rag/punkt.zip\n",
    "!wget https://modelscope.oss-cn-beijing.aliyuncs.com/resource/rag/stopwords.zip\n",
    "!wget https://modelscope.oss-cn-beijing.aliyuncs.com/resource/rag/xianjiaoda.md\n",
    "\n",
    "!mkdir -p /opt/conda/lib/python3.10/site-packages/llama_index/core/_static/nltk_cache/tokenizers\n",
    "!mkdir -p /opt/conda/lib/python3.10/site-packages/llama_index/core/_static/nltk_cache/corpora\n",
    "\n",
    "!cp /mnt/workspace/punkt.zip /opt/conda/lib/python3.10/site-packages/llama_index/core/_static/nltk_cache/tokenizers\n",
    "!cp /mnt/workspace/stopwords.zip /opt/conda/lib/python3.10/site-packages/llama_index/core/_static/nltk_cache/corpora\n",
    "!cd /opt/conda/lib/python3.10/site-packages/llama_index/core/_static/nltk_cache/tokenizers; unzip punkt.zip;\n",
    "!cd /opt/conda/lib/python3.10/site-packages/llama_index/core/_static/nltk_cache/corpora; unzip stopwords.zip;\n",
    "\n",
    "\n",
    "!mkdir -p /mnt/workspace/custom_data\n",
    "!mv /mnt/workspace/xianjiaoda.md /mnt/workspace/custom_data\n",
    "\n",
    "!cd /mnt/workspace"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "63704e2b21a9ba52"
  },
  {
   "cell_type": "code",
   "outputs": [],
   "source": [
    "import logging\n",
    "import sys\n",
    "from abc import ABC\n",
    "from typing import Any, List\n",
    "\n",
    "import torch\n",
    "from llama_index.core import (\n",
    "    SimpleDirectoryReader,\n",
    "    VectorStoreIndex,\n",
    "    Settings,\n",
    "    ServiceContext,\n",
    "    set_global_service_context,\n",
    ")\n",
    "from llama_index.core.base.embeddings.base import BaseEmbedding, Embedding\n",
    "from llama_index.core.prompts import PromptTemplate\n",
    "from llama_index.llms.huggingface import HuggingFaceLLM\n",
    "\n",
    "from modelscope import snapshot_download\n",
    "\n",
    "logging.basicConfig(stream=sys.stdout, level=logging.INFO)\n",
    "logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))\n",
    "\n",
    "# download QWEN model from modelscope\n",
    "qwen15_4B_CHAT = \"qwen/Qwen1.5-4B-Chat\"\n",
    "selected_model = snapshot_download(qwen15_4B_CHAT)\n",
    "\n",
    "# define sys prompt\n",
    "SYSTEM_PROMPT = \"\"\"You are a helpful AI assistant.\"\"\"\n",
    "query_wrapper_prompt = PromptTemplate(\n",
    "    \"[INST]<<SYS>>\\n\" + SYSTEM_PROMPT + \"<</SYS>>\\n\\n{query_str}[/INST] \"\n",
    ")\n",
    "\n",
    "# create HuggingFaceLLM with qwen1.5 \n",
    "llm = HuggingFaceLLM(\n",
    "    context_window=4096,\n",
    "    max_new_tokens=2048,\n",
    "    generate_kwargs={\"temperature\": 0.0, \"do_sample\": False},\n",
    "    query_wrapper_prompt=query_wrapper_prompt,\n",
    "    tokenizer_name=selected_model,\n",
    "    model_name=selected_model,\n",
    "    device_map=\"auto\",\n",
    "    # change these settings below depending on your GPU\n",
    "    model_kwargs={\"torch_dtype\": torch.float16},\n",
    ")\n",
    "print(\"llm created\")\n",
    "\n",
    "\n",
    "# wrap modelscope embedding for llama-index (based on BaseEmbedding)\n",
    "class ModelScopeEmbeddings4LlamaIndex(BaseEmbedding, ABC):\n",
    "    embed: Any = None\n",
    "    model_id: str = \"damo/nlp_gte_sentence-embedding_chinese-base\"\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",
    "            # 使用modelscope的embedding模型（包含下载）\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) -> Embedding:\n",
    "        text = query.replace(\"\\n\", \" \")\n",
    "        inputs = {\"source_sentence\": [text]}\n",
    "        # note that we have to call tolist() to change numpy.ndarray into python list\n",
    "        return self.embed(input=inputs)['text_embedding'][0].tolist()\n",
    "\n",
    "    def _get_text_embedding(self, text: str) -> Embedding:\n",
    "        text = text.replace(\"\\n\", \" \")\n",
    "        inputs = {\"source_sentence\": [text]}\n",
    "        return self.embed(input=inputs)['text_embedding'][0].tolist()\n",
    "\n",
    "    def _get_text_embeddings(self, texts: List[str]) -> List[Embedding]:\n",
    "        texts = list(map(lambda x: x.replace(\"\\n\", \" \"), texts))\n",
    "        inputs = {\"source_sentence\": texts}\n",
    "        return self.embed(input=inputs)['text_embedding'].tolist()\n",
    "\n",
    "    async def _aget_query_embedding(self, query: str) -> Embedding:\n",
    "        return self._get_query_embedding(query)\n",
    "\n",
    "\n",
    "embedding_model = \"damo/nlp_gte_sentence-embedding_chinese-base\"\n",
    "embeddings = ModelScopeEmbeddings4LlamaIndex(model_id=embedding_model)\n",
    "service_context = ServiceContext.from_defaults(embed_model=embeddings, llm=llm)\n",
    "set_global_service_context(service_context)\n",
    "Settings.embed_model = embeddings\n",
    "\n",
    "# load example documents\n",
    "documents = SimpleDirectoryReader(\"/mnt/workspace/custom_data/\").load_data()\n",
    "\n",
    "# create Vector DB\n",
    "index = VectorStoreIndex.from_documents(documents)\n",
    "\n",
    "# set Logging to DEBUG for more detailed outputs\n",
    "query_engine = index.as_query_engine()\n",
    "\n",
    "# do query\n",
    "response = query_engine.query(\"西安较大的校训是什么\")\n",
    "print(response)\n"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "eef67659e94045c5"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
