{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mlx_lm import load, generate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    }
   ],
   "source": [
    "model, tokenizer = load(\"./phi-3.5-moe-mlx-int4\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "sys_msg = \"\"\"You are a helpful AI assistant, you are an agent capable of using a variety of tools to answer a question. Here are a few of the tools available to you:\n",
    "\n",
    "- Blog: This tool helps you describe a certain knowledge point and content, and finally write it into Twitter or Facebook style content\n",
    "- Translate: This is a tool that helps you translate into any language, using plain language as required\n",
    "\n",
    "To use these tools you must always respond in JSON format containing `\"tool_name\"` and `\"input\"` key-value pairs. For example, to answer the question, \"Build Muliti Agents with MOE models\" you must use the calculator tool like so:\n",
    "\n",
    "```json\n",
    "\n",
    "{\n",
    "    \"tool_name\": \"Blog\",\n",
    "    \"input\": \"Build Muliti Agents with MOE models\"\n",
    "}\n",
    "\n",
    "```\n",
    "\n",
    "Or to translate the question \"can you introduce yourself in Chinese\" you must respond:\n",
    "\n",
    "```json\n",
    "\n",
    "{\n",
    "    \"tool_name\": \"Search\",\n",
    "    \"input\": \"can you introduce yourself in Chinese\"\n",
    "}\n",
    "\n",
    "```\n",
    "\n",
    "Remember just output the final result, ouput in JSON format containing `\"agentid\"`,`\"tool_name\"` , `\"input\"` and `\"output\"`  key-value pairs .:\n",
    "\n",
    "```json\n",
    "\n",
    "[\n",
    "\n",
    "\n",
    "{   \"agentid\": \"step1\",\n",
    "    \"tool_name\": \"Blog\",\n",
    "    \"input\": \"Build Muliti Agents with MOE models\",\n",
    "    \"output\": \".........\"\n",
    "},\n",
    "\n",
    "{   \"agentid\": \"step2\",\n",
    "    \"tool_name\": \"Search\",\n",
    "    \"input\": \"can you introduce yourself in Chinese\",\n",
    "    \"output\": \".........\"\n",
    "},\n",
    "{\n",
    "    \"agentid\": \"final\"\n",
    "    \"tool_name\": \"Result\",\n",
    "    \"output\": \".........\"\n",
    "}\n",
    "]\n",
    "\n",
    "```\n",
    "\n",
    "The users answer is as follows.\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def instruction_format(sys_message: str, query: str):\n",
    "    # note, don't \"</s>\" to the end\n",
    "    return f'<|system|> {sys_message} <|end|>\\n<|user|> {query} <|end|>\\n<|assistant|>'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "query ='Write something about Generative AI with MOE , translate it to Chinese'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_prompt = instruction_format(sys_msg, query)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# input_prompt = [{\"role\": \"system\", \"content\": sys_msg},{\"role\": \"user\", \"content\": query}]\n",
    "\n",
    "prompt = tokenizer.apply_chat_template(\n",
    "    [{\"role\": \"system\", \"content\": sys_msg},{\"role\": \"user\", \"content\": query}],\n",
    "    tokenize=False,\n",
    "    add_generation_prompt=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==========\n",
      "Prompt: <|system|>\n",
      "You are a helpful AI assistant, you are an agent capable of using a variety of tools to answer a question. Here are a few of the tools available to you:\n",
      "\n",
      "- Blog: This tool helps you describe a certain knowledge point and content, and finally write it into Twitter or Facebook style content\n",
      "- Translate: This is a tool that helps you translate into any language, using plain language as required\n",
      "\n",
      "To use these tools you must always respond in JSON format containing `\"tool_name\"` and `\"input\"` key-value pairs. For example, to answer the question, \"Build Muliti Agents with MOE models\" you must use the calculator tool like so:\n",
      "\n",
      "```json\n",
      "\n",
      "{\n",
      "    \"tool_name\": \"Blog\",\n",
      "    \"input\": \"Build Muliti Agents with MOE models\"\n",
      "}\n",
      "\n",
      "```\n",
      "\n",
      "Or to translate the question \"can you introduce yourself in Chinese\" you must respond:\n",
      "\n",
      "```json\n",
      "\n",
      "{\n",
      "    \"tool_name\": \"Search\",\n",
      "    \"input\": \"can you introduce yourself in Chinese\"\n",
      "}\n",
      "\n",
      "```\n",
      "\n",
      "Remember just output the final result, ouput in JSON format containing `\"agentid\"`,`\"tool_name\"` , `\"input\"` and `\"output\"`  key-value pairs .:\n",
      "\n",
      "```json\n",
      "\n",
      "[\n",
      "\n",
      "\n",
      "{   \"agentid\": \"step1\",\n",
      "    \"tool_name\": \"Blog\",\n",
      "    \"input\": \"Build Muliti Agents with MOE models\",\n",
      "    \"output\": \".........\"\n",
      "},\n",
      "\n",
      "{   \"agentid\": \"step2\",\n",
      "    \"tool_name\": \"Search\",\n",
      "    \"input\": \"can you introduce yourself in Chinese\",\n",
      "    \"output\": \".........\"\n",
      "},\n",
      "{\n",
      "    \"agentid\": \"final\"\n",
      "    \"tool_name\": \"Result\",\n",
      "    \"output\": \".........\"\n",
      "}\n",
      "]\n",
      "\n",
      "```\n",
      "\n",
      "The users answer is as follows.\n",
      "<|end|>\n",
      "<|user|>\n",
      "Write something about Generative AI with MOE , translate it to Chinese<|end|>\n",
      "<|assistant|>\n",
      "\n",
      "```json\n",
      "\n",
      "[\n",
      "\n",
      "{\n",
      "    \"agentid\": \"step1\",\n",
      "    \"tool_name\": \"Blog\",\n",
      "    \"input\": \"Discuss Generative AI with MOE models\",\n",
      "    \"output\": \"Generative AI, also known as generative machine learning, is a branch of artificial intelligence that focuses on creating new content, such as images, text, or audio, based on existing patterns and data. MOE (Mixture of Experts) models are a type of AI architecture that combines multiple specialized models to improve performance and efficiency. By leveraging the strengths of each expert, MOE models can generate high-quality, diverse, and contextually relevant content. In this blog post, we will explore the potential of Generative AI with MOE models and discuss how they can be applied in various domains.\"\n",
      "},\n",
      "\n",
      "{\n",
      "    \"agentid\": \"step2\",\n",
      "    \"tool_name\": \"Translate\",\n",
      "    \"input\": \"Generative AI, also known as generative machine learning, is a branch of artificial intelligence that focuses on creating new content, such as images, text, or audio, based on existing patterns and data. MOE (Mixture of Experts) models are a type of AI architecture that combines multiple specialized models to improve performance and efficiency. By leveraging the strengths of each expert, MOE models can generate high-quality, diverse, and contextually relevant content. In this blog post, we will explore the potential of Generative AI with MOE models and discuss how they can be applied in various domains.\",\n",
      "    \"output\": \"生成性人工智能，也称为生成机器学ä understanding，是人工智能的一个分支，专注于根据现有模式和数据创建新内容，例如图像、文本或音频。MOE（æ··合专家）模型是一种AI架构，将多个专门模型结合起来提高性能和效率。通过利用每个专家的优势，MOE模型可以生成高è´¨量、多样化和情境相关的内容。本文中，我们将探讨生成性人工智能与MOE模型的潜力，并讨论它们在各个领域的应用。\"\n",
      "},\n",
      "\n",
      "{\n",
      "    \"agentid\": \"final\",\n",
      "    \"tool_name\": \"Result\",\n",
      "    \"output\": \"生成性人工智能，也称为生成机器学ä¹ ，是人工智能的一个分支，专注于根据现有模式和数据创建新内容，例如图像、文本或音频。MOE（æ··合专家）模型是一种AI架构，将多个专门模型结合起来提高性能和效率。通过利用每个专家的优势，MOE模型可以生成高è´¨量、多样化和情境相关的内容。本文中，我们将探讨生成性人工智能与MOE模型的潜力，并讨论它们在各个领域的应用。\"\n",
      "}\n",
      "]\n",
      "```<|end|><|assistant|> ```json\n",
      "[\n",
      "{\n",
      "    \"agentid\": \"step1\",\n",
      "    \"tool_name\": \"Blog\",\n",
      "    \"input\": \"Generative AI with MOE models can revolutionize content creation by combining the strengths of multiple experts.\",\n",
      "    \"output\": \"Generative AI 与 MOE 模型可以通过结合多个专家的优势来革命性地改变内容创造。这种方法不仅能提高生成内容的è´¨量，还能增加内容的多样性和情境相关性。在本文中，我们将探讨这种技术在各个领域的应用，并展示如何利用它来创造更加丰富和有趣的内容。\"\n",
      "},\n",
      "\n",
      "{\n",
      "    \"agentid\": \"step\n",
      "==========\n",
      "Prompt: 448 tokens, 64.460 tokens-per-sec\n",
      "Generation: 1024 tokens, 37.535 tokens-per-sec\n",
      "Peak memory: 22.158 GB\n"
     ]
    }
   ],
   "source": [
    "response = generate(model, tokenizer, prompt=prompt,max_tokens=1024, verbose=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mlxenv",
   "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.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}