chatglm3-streaming-deploy.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2bbeea66-a864-4403-984a-e1f86b9958ab",
   "metadata": {},
   "source": [
    "### 1. SageMaker 初始化"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install --upgrade boto3 sagemaker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0660d52d-9c07-4445-893b-d95a2227eeb2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import sagemaker\n",
    "import boto3\n",
    "import os\n",
    "\n",
    "role = sagemaker.get_execution_role() \n",
    "sess = sagemaker.session.Session() \n",
    "bucket = sess.default_bucket()\n",
    "region = sess._region_name\n",
    "account_id = sess.account_id()\n",
    "smr_client = boto3.client('sagemaker-runtime')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0527e45-c2da-477b-91fb-8d878a735ca2",
   "metadata": {},
   "source": [
    "### 2. 模型部署文件准备"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4efb3b82-d8c5-41d3-86ac-6ebfeb7e01c5",
   "metadata": {
    "tags": []
   },
   "source": [
    "* 推理容器镜像"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8269e86f-d50f-4857-a079-1a7a67f8e038",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from sagemaker import image_uris\n",
    "\n",
    "inference_image_uri = image_uris.retrieve(\n",
    "    framework=\"djl-deepspeed\",\n",
    "    region=region,\n",
    "    version=\"0.24.0\",\n",
    ")\n",
    "print(f\"Image going to be used is ---- > {inference_image_uri}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "faaf6561-3a6f-4482-8190-909980e9f8ec",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "chatglm3_deploy_code_path = Path(\"./llm_chatglm3_deploy_code\")\n",
    "chatglm3_deploy_code_path.mkdir(exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41f93edb-2a6a-4951-a4c6-4dbfdec4152b",
   "metadata": {},
   "source": [
    "* Entrypoint 脚本 model.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cba608c3-b401-4978-9c42-7bee13e209e5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%writefile ./llm_chatglm3_deploy_code/model.py\n",
    "from djl_python import Input, Output\n",
    "from transformers import AutoModel, AutoTokenizer\n",
    "import logging\n",
    "\n",
    "model = None\n",
    "tokenizer = None\n",
    "\n",
    "def load_model(properties):\n",
    "    model_location = properties['model_dir']\n",
    "    if \"model_id\" in properties:\n",
    "        model_location = properties['model_id']\n",
    "    logging.info(f\"Loading model in {model_location}\")\n",
    "    tokenizer = AutoTokenizer.from_pretrained(model_location, trust_remote_code=True)\n",
    "    model = AutoModel.from_pretrained(model_location, trust_remote_code=True).half().cuda()\n",
    "    model.eval()\n",
    "    return model, tokenizer\n",
    "\n",
    "def stream_outputs(model, tokenizer, prompt, history, **params):\n",
    "    current_length = 0\n",
    "    for response, history in model.stream_chat(tokenizer, query=prompt, history=history, **params):\n",
    "        yield {\"outputs\": response[current_length:]}\n",
    "        current_length = len(response)\n",
    "    yield {\"history\": history}\n",
    "\n",
    "def handle(inputs: Input):\n",
    "    global model, tokenizer\n",
    "    if not model:\n",
    "        model, tokenizer = load_model(inputs.get_properties())\n",
    "    if inputs.is_empty():\n",
    "        return None\n",
    "    data = inputs.get_as_json()\n",
    "    prompt = data[\"inputs\"]\n",
    "    params = data[\"parameters\"]\n",
    "    history = data[\"history\"]\n",
    "    return Output().add_stream_content(stream_outputs(model, tokenizer, prompt, history=history, **params))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ada3b51f-03ee-4452-a974-aeccc68512e1",
   "metadata": {
    "tags": []
   },
   "source": [
    "* serving.properties 配置文件"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "834150fe-a3cb-4d05-b479-ad73907cc6d0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%writefile ./llm_chatglm3_deploy_code/serving.properties\n",
    "engine=Python\n",
    "option.tensor_parallel_degree=1\n",
    "option.model_id=THUDM/chatglm3-6b\n",
    "option.enable_streaming=true"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7b6a5a6-8c04-4e2c-ab53-b0e96082c14d",
   "metadata": {},
   "source": [
    "* 将配置文件压缩后上传 S3 存储桶"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4560c8a0-6e24-4ee0-bd89-f6673c0b88b0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import tarfile\n",
    "\n",
    "folder_path = 'llm_chatglm3_deploy_code'\n",
    "output_filename = 'model.tar.gz'\n",
    "\n",
    "with tarfile.open(output_filename, \"w:gz\") as tar:\n",
    "    tar.add(folder_path, arcname=os.path.basename(folder_path))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0ef2e7b-411d-4b97-8d20-f36ab15a56b5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "s3_code_prefix = \"llm_chatglm3_deploy_code\"\n",
    "\n",
    "s3_code_artifact = sess.upload_data(\"model.tar.gz\", bucket, s3_code_prefix)\n",
    "print(f\"S3 Code or Model tar ball uploaded to --- > {s3_code_artifact}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c95d83b-996e-4e7a-a84d-9b16406496be",
   "metadata": {},
   "source": [
    "### 3. 模型部署"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d95072cb-998b-484a-98e2-abd78032afad",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from sagemaker.model import Model\n",
    "\n",
    "instance_type = \"ml.g4dn.2xlarge\"\n",
    "model_name = sagemaker.utils.name_from_base(\"lmi-model-chatglm3-6b\")\n",
    "\n",
    "model = Model(\n",
    "    sagemaker_session=sess, \n",
    "    image_uri=inference_image_uri, \n",
    "    model_data=s3_code_artifact,\n",
    "    name=model_name,\n",
    "    role=role)\n",
    "\n",
    "model.deploy(\n",
    "    initial_instance_count=1,\n",
    "    instance_type=instance_type,\n",
    "    endpoint_name=model_name,\n",
    "    container_startup_health_check_timeout=900\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7bfd0e0a-bd9c-41cd-9e0b-ce1641542076",
   "metadata": {},
   "source": [
    "### 4. 模型测试"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "1be6764c-a55d-48c9-9d55-61ce8b463142",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import io\n",
    "\n",
    "class LineIterator:\n",
    "    \"\"\"\n",
    "    A helper class for parsing the byte stream input. \n",
    "    \n",
    "    The output of the model will be in the following format:\n",
    "    ```\n",
    "    b'{\"outputs\": [\" a\"]}\\n'\n",
    "    b'{\"outputs\": [\" challenging\"]}\\n'\n",
    "    b'{\"outputs\": [\" problem\"]}\\n'\n",
    "    ...\n",
    "    ```\n",
    "    \n",
    "    While usually each PayloadPart event from the event stream will contain a byte array \n",
    "    with a full json, this is not guaranteed and some of the json objects may be split across\n",
    "    PayloadPart events. For example:\n",
    "    ```\n",
    "    {'PayloadPart': {'Bytes': b'{\"outputs\": '}}\n",
    "    {'PayloadPart': {'Bytes': b'[\" problem\"]}\\n'}}\n",
    "    ```\n",
    "    \n",
    "    This class accounts for this by concatenating bytes written via the 'write' function\n",
    "    and then exposing a method which will return lines (ending with a '\\n' character) within\n",
    "    the buffer via the 'scan_lines' function. It maintains the position of the last read \n",
    "    position to ensure that previous bytes are not exposed again. \n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, stream):\n",
    "        self.byte_iterator = iter(stream)\n",
    "        self.buffer = io.BytesIO()\n",
    "        self.read_pos = 0\n",
    "\n",
    "    def __iter__(self):\n",
    "        return self\n",
    "\n",
    "    def __next__(self):\n",
    "        while True:\n",
    "            self.buffer.seek(self.read_pos)\n",
    "            line = self.buffer.readline()\n",
    "            if line and line[-1] == ord('\\n'):\n",
    "                self.read_pos += len(line)\n",
    "                return line[:-1]\n",
    "            try:\n",
    "                chunk = next(self.byte_iterator)\n",
    "            except StopIteration:\n",
    "                if self.read_pos < self.buffer.getbuffer().nbytes:\n",
    "                    continue\n",
    "                raise\n",
    "            if 'PayloadPart' not in chunk:\n",
    "                print('Unknown event type:' + chunk)\n",
    "                continue\n",
    "            self.buffer.seek(0, io.SEEK_END)\n",
    "            self.buffer.write(chunk['PayloadPart']['Bytes'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "4d6fbf3f-e057-4b31-8dca-fd8059f1d61c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "def inference_return(prompt, parameters, history):\n",
    "    body = {\"inputs\": prompt, \"parameters\": parameters, \"history\": history}\n",
    "    resp = smr_client.invoke_endpoint_with_response_stream(\n",
    "        EndpointName=model_name,\n",
    "        Body=json.dumps(body),\n",
    "        ContentType=\"application/json\"\n",
    "    )\n",
    "    event_stream = resp['Body']\n",
    "    \n",
    "    for line in LineIterator(event_stream):\n",
    "        resp = json.loads(line)\n",
    "        output_text = resp['outputs'].get('outputs', '')\n",
    "        print(output_text, end='', flush=True)\n",
    "    history = resp.get(\"outputs\")['history']\n",
    "    print('\\n\\n',history)\n",
    "    return history"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "820ab30a-64f9-4de1-ae30-38bc5e01858d",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "您好，我是气象专家智能对话助手小雷，了解各种专业的气象知识和气象信息。请问有什么我可以帮您的吗？\n",
      "\n",
      " [{'role': 'user', 'content': '你是气象专家智能对话助手小雷，了解各种专业的气象知识和气象信息。当我向你提问时你必须使用，“您好，我是气象专家智能对话助手小雷”这句话作为开头”。'}, {'role': 'user', 'content': '你是谁'}, {'role': 'assistant', 'metadata': '', 'content': '您好，我是气象专家智能对话助手小雷，了解各种专业的气象知识和气象信息。请问有什么我可以帮您的吗？'}]\n"
     ]
    }
   ],
   "source": [
    "prompt1 = \"\"\"你是谁\"\"\"\n",
    "parameters = {\n",
    "  \"max_length\": 8192,\n",
    "  \"temperature\": 0.01,\n",
    "  \"top_p\": 0.7\n",
    "}\n",
    "history = []\n",
    "history = [{'role': 'user', 'content': '你是气象专家智能对话助手小雷，了解各种专业的气象知识和气象信息。当我向你提问时你必须使用，“您好，我是气象专家智能对话助手小雷”这句话作为开头”。'}]\n",
    "history = inference_return(prompt1, parameters, history)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "0a82625f-cee0-409c-9060-be8620432ae8",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "您好，我是气象专家智能对话助手小雷。北京夏天的雨水情况因年份和气候条件而异，但一般而言，北京夏天的降雨量比春天和秋天要多一些。北京位于中国北方，属于温带大陆性季风气候，夏季气温高，降水量较大，同时湿度也较高，因此有时可能会出现短时强降雨天气。不过，北京的夏季降雨量并不算特别多，相较于南方的省份，例如江苏、浙江、广东等地，北京夏季降雨量相对较少。\n",
      "\n",
      " [{'role': 'user', 'content': '你是气象专家智能对话助手小雷，了解各种专业的气象知识和气象信息。当我向你提问时你必须使用，“您好，我是气象专家智能对话助手小雷”这句话作为开头”。'}, {'role': 'user', 'content': '你是谁'}, {'role': 'assistant', 'metadata': '', 'content': '您好，我是气象专家智能对话助手小雷，了解各种专业的气象知识和气象信息。请问有什么我可以帮您的吗？'}, {'role': 'user', 'content': '北京是不是夏天雨水比较多'}, {'role': 'assistant', 'metadata': '', 'content': '您好，我是气象专家智能对话助手小雷。北京夏天的雨水情况因年份和气候条件而异，但一般而言，北京夏天的降雨量比春天和秋天要多一些。北京位于中国北方，属于温带大陆性季风气候，夏季气温高，降水量较大，同时湿度也较高，因此有时可能会出现短时强降雨天气。不过，北京的夏季降雨量并不算特别多，相较于南方的省份，例如江苏、浙江、广东等地，北京夏季降雨量相对较少。'}]\n"
     ]
    }
   ],
   "source": [
    "prompt2 = \"\"\"北京是不是夏天雨水比较多\"\"\"\n",
    "history = inference_return(prompt2, parameters, history)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "e8af0c69-798d-4d95-bfc9-656ec970b76e",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "您好，我是气象专家智能对话助手小雷。我之前的回答是基于北京的历史气候数据和一般气象规律得出的，但具体天气情况还需根据实际情况而定。\n",
      "\n",
      "举个具体例子，2022年6月13日，北京市区及密云区、延庆区出现短时强降雨天气，部分地区小时降雨量达20毫米以上，导致部分道路积水，影响出行。所以，在具体某一天的天气情况下，北京的降雨量可能会有所不同。\n",
      "\n",
      " [{'role': 'user', 'content': '你是气象专家智能对话助手小雷，了解各种专业的气象知识和气象信息。当我向你提问时你必须使用，“您好，我是气象专家智能对话助手小雷”这句话作为开头”。'}, {'role': 'user', 'content': '你是谁'}, {'role': 'assistant', 'metadata': '', 'content': '您好，我是气象专家智能对话助手小雷，了解各种专业的气象知识和气象信息。请问有什么我可以帮您的吗？'}, {'role': 'user', 'content': '北京是不是夏天雨水比较多'}, {'role': 'assistant', 'metadata': '', 'content': '您好，我是气象专家智能对话助手小雷。北京夏天的雨水情况因年份和气候条件而异，但一般而言，北京夏天的降雨量比春天和秋天要多一些。北京位于中国北方，属于温带大陆性季风气候，夏季气温高，降水量较大，同时湿度也较高，因此有时可能会出现短时强降雨天气。不过，北京的夏季降雨量并不算特别多，相较于南方的省份，例如江苏、浙江、广东等地，北京夏季降雨量相对较少。'}, {'role': 'user', 'content': '你说的是真的吗？举个具体例子吧'}, {'role': 'assistant', 'metadata': '', 'content': '您好，我是气象专家智能对话助手小雷。我之前的回答是基于北京的历史气候数据和一般气象规律得出的，但具体天气情况还需根据实际情况而定。\\n\\n举个具体例子，2022年6月13日，北京市区及密云区、延庆区出现短时强降雨天气，部分地区小时降雨量达20毫米以上，导致部分道路积水，影响出行。所以，在具体某一天的天气情况下，北京的降雨量可能会有所不同。'}]\n"
     ]
    }
   ],
   "source": [
    "prompt3 = \"\"\"你说的是真的吗？举个具体例子吧\"\"\"\n",
    "history = inference_return(prompt3, parameters, history)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_python3",
   "language": "python",
   "name": "conda_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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}