{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "285db456",
   "metadata": {},
   "source": [
    "### SageMaker  Endpoint 部署ChatYuan\n",
    "  \n",
    "[ChatYuan](https://huggingface.co/ClueAI/ChatYuan-large-v1): 元语功能型对话大模型 , 这个模型可以用于问答、结合上下文做对话、做各种生成任务，包括创意性写作，也能回答一些像法律、新冠等领域问题。它基于PromptCLUE-large结合数亿条功能对话多轮对话数据进一步训练得到。\n",
    "\n",
    "PromptCLUE-large在1000亿token中文语料上预训练，累计学习1.5万亿中文token，并且在数百种任务上进行Prompt任务式训练。针对理解类任务，如分类、情感分析、抽取等，可以自定义标签体系；针对多种生成任务，可以进行采样自由生成。\n",
    "\n",
    "#### 准备\n",
    "1. 升级boto3, sagemaker python sdk  \n",
    "2. 准备inference.py, requirements.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3db7cd5",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install --upgrade boto3\n",
    "!pip install --upgrade sagemaker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad9c4f2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "import sagemaker\n",
    "\n",
    "account_id = boto3.client('sts').get_caller_identity().get('Account')\n",
    "region_name = boto3.session.Session().region_name\n",
    "\n",
    "sagemaker_session = sagemaker.Session()\n",
    "bucket = sagemaker_session.default_bucket()\n",
    "role = sagemaker.get_execution_role()\n",
    "\n",
    "print(role)\n",
    "print(bucket)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c01f268f",
   "metadata": {},
   "outputs": [],
   "source": [
    "!touch dummy\n",
    "!tar czvf model.tar.gz dummy\n",
    "assets_dir = 's3://{0}/{1}/assets/'.format(bucket, 'chatyuan')\n",
    "model_data = 's3://{0}/{1}/assets/model.tar.gz'.format(bucket, 'chatyuan')\n",
    "!aws s3 cp model.tar.gz $assets_dir\n",
    "!rm -f dummy model.tar.gz"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52cf6808",
   "metadata": {},
   "source": [
    "### 设置模型推理参数\n",
    "1. model_name: Huggingface diffusers models (not support single check point format)\n",
    "2. model_args: diffuser StableDiffusionPipeline init arguments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed0d7adf",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = None\n",
    "entry_point = 'inference.py'\n",
    "framework_version = '1.12'\n",
    "py_version = 'py38'\n",
    "bucket='m.flowq.io'\n",
    "model_environment = {\n",
    "    'SAGEMAKER_MODEL_SERVER_TIMEOUT':'600', \n",
    "    'SAGEMAKER_MODEL_SERVER_WORKERS': '1', \n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c641854e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.pytorch.model import PyTorchModel\n",
    "\n",
    "model = PyTorchModel(\n",
    "    name = model_name,\n",
    "    model_data = model_data,\n",
    "    entry_point = 'inference.py',\n",
    "    source_dir = './code',\n",
    "    role = role,\n",
    "    framework_version = framework_version, \n",
    "    py_version = py_version,\n",
    "    env = model_environment\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2eaedc75",
   "metadata": {},
   "source": [
    "### 部署 SageMaker Endpoint\n",
    "部署推理服务"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74a9b7e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "endpoint_name = None\n",
    "instance_type = 'ml.g4dn.xlarge'\n",
    "instance_count = 1\n",
    "\n",
    "from sagemaker.serializers import JSONSerializer\n",
    "from sagemaker.deserializers import JSONDeserializer\n",
    "predictor = model.deploy(\n",
    "    endpoint_name = endpoint_name,\n",
    "    instance_type = instance_type, \n",
    "    initial_instance_count = instance_count,\n",
    "    serializer = JSONSerializer(),\n",
    "    deserializer = JSONDeserializer()\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25218d7d",
   "metadata": {},
   "source": [
    "### ChatYuan 测试\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "487ee8e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs= {\n",
    "    \"ask\": \"费德勒网球水平怎么样? \"\n",
    "\n",
    "}\n",
    "\n",
    "response = predictor.predict(inputs)\n",
    "print(response[\"answer\"])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f75bb18",
   "metadata": {},
   "outputs": [],
   "source": [
    "#我们来查看一下刚部署的这个chatyuan模型对应的SageMaker inference endpoint\n",
    "your_chatyuan_endpoint_name = predictor.endpoint_name\n",
    "your_chatyuan_endpoint_name"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "463d21ef",
   "metadata": {},
   "source": [
    "利用已经在SageMaker real time inference endpoint部署的ChatYuan LLM模型来模拟单轮对话和多轮对话。\n",
    "\n",
    "下面的代码建议在SageMaker Notebook上来运行。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "239e990a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import boto3\n",
    "\n",
    "client = boto3.client('runtime.sagemaker')\n",
    "\n",
    "def query_endpoint_with_json_payload(encoded_json):\n",
    "    response = client.invoke_endpoint(EndpointName=your_chatyuan_endpoint_name, ContentType='application/json', Body=encoded_json)\n",
    "    return response\n",
    "\n",
    "def parse_response_texts(query_response):\n",
    "    model_predictions = json.loads(query_response['Body'].read())\n",
    "    generated_text = model_predictions[\"answer\"]\n",
    "    return generated_text\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1763ebc0",
   "metadata": {},
   "source": [
    "先简单测试一下chatyuan针对单个问题的回答"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63d9d838",
   "metadata": {},
   "outputs": [],
   "source": [
    "payload = {\"ask\": \"信息抽取：\\n2022年世界杯的冠军是阿根廷队伍，梅西是MVP\\n问题：国家名，人名\\n答案：\"}\n",
    "query_response = query_endpoint_with_json_payload(json.dumps(payload).encode('utf-8'))\n",
    "generated_texts = parse_response_texts(query_response)\n",
    "print(generated_texts)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e48a2fe1",
   "metadata": {},
   "source": [
    "单轮对话：每个问题/query都是独立的，问题之间没有关联性。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1450b79",
   "metadata": {},
   "outputs": [],
   "source": [
    "#1.首先需要一个简单的开场拍。\n",
    "print(\"用户：你好！\\n小元：您好!我是元语AI。我可以回答您的问题、写文章、写作业、翻译，对于一些法律等领域的问题我也可以给你提供信息。\")\n",
    "\n",
    "#2.在同一个session中持续对话，但是每次对话之间没有什么关联。\n",
    "while True:\n",
    "    command = input(\"用户：\")\n",
    "    if command == 'quit':\n",
    "        break;\n",
    "    \n",
    "    #print(command)\n",
    "    payload = {\"ask\": \"\\n用户：\"+ command + \"\\n\"}\n",
    "    #print(payload[\"ask\"])\n",
    "    query_response = query_endpoint_with_json_payload(json.dumps(payload).encode('utf-8'))\n",
    "    generated_texts = \"小元：\" + parse_response_texts(query_response)\n",
    "    print(generated_texts)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2ac2370",
   "metadata": {},
   "source": [
    "多轮对话模拟：我们这里来测试一下chatyuan的多轮对话能力。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acdfd4b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "#1.首先需要开场拍来引导chatyuan，其实就是给它一个上下文来启动所谓的对话session。\n",
    "payload = {\"ask\": \"用户：你好！\\n小元：您好!我是元语AI。我可以回答您的问题、写文章、写作业、翻译，对于一些法律等领域的问题我也可以给你提供信息。\"}\n",
    "print(payload[\"ask\"])\n",
    "generated_texts = \"\"\n",
    "\n",
    "#在这里简单模拟多轮对话时，发送给SageMaker endpoint的payload会越来越大，这里对payload大致做一个限制。\n",
    "session_len = 10 * 1024 * 1024 \n",
    "\n",
    "#2.在同一个session中持续对话，为了有多轮对话的效果，把每一轮的信息(问题-回答对)都带上来告诉chatyuan之前的上下文。\n",
    "while True:\n",
    "    command = input(\"用户：\")\n",
    "    if command == 'quit':\n",
    "        break;\n",
    "    \n",
    "    #print(command)\n",
    "    whole_context = payload[\"ask\"] + generated_texts + \"\\n用户：\"+ command + \"\\n\"\n",
    "    payload = {\"ask\": whole_context}\n",
    "    if len(whole_context) > session_len:\n",
    "        print(\"上下文信息太多了，当前对话session要退出了！\")\n",
    "        break;\n",
    "    #print(payload[\"ask\"])\n",
    "    query_response = query_endpoint_with_json_payload(json.dumps(payload).encode('utf-8'))\n",
    "    generated_texts = \"小元：\" + parse_response_texts(query_response)\n",
    "    print(generated_texts)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4aff6ebe",
   "metadata": {},
   "source": [
    "### 删除SageMaker  Endpoint\n",
    "删除推理服务"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0938bc68",
   "metadata": {},
   "outputs": [],
   "source": [
    "predictor.delete_endpoint()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98078bd9",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.6.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}