{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction to Conversational AI with MLflow and DialoGPT\n",
    "\n",
    "Welcome to our tutorial on integrating [Microsoft's DialoGPT](https://huggingface.co/microsoft/DialoGPT-medium) with MLflow's transformers flavor to explore conversational AI."
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "<a href=\"https://raw.githubusercontent.com/mlflow/mlflow/master/docs/source/llms/transformers/tutorials/conversational/conversational-model.ipynb\" class=\"notebook-download-btn\">Download this Notebook</a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "### Learning Objectives\n",
    "\n",
    "In this tutorial, you will:\n",
    "\n",
    "- Set up a conversational AI **pipeline** using DialoGPT from the Transformers library.\n",
    "- **Log** the DialoGPT model along with its configurations using MLflow.\n",
    "- Infer the input and output **signature** of the DialoGPT model.\n",
    "- **Load** a stored DialoGPT model from MLflow for interactive usage.\n",
    "- Interact with the chatbot model and understand the nuances of conversational AI.\n",
    "\n",
    "By the end of this tutorial, you will have a solid understanding of managing and deploying conversational AI models with MLflow, enhancing your capabilities in natural language processing.\n",
    "\n",
    "#### What is DialoGPT?\n",
    "DialoGPT is a conversational model developed by Microsoft, fine-tuned on a large dataset of dialogues to generate human-like responses. Part of the GPT family, DialoGPT excels in natural language understanding and generation, making it ideal for chatbots.\n",
    "\n",
    "#### Why MLflow with DialoGPT?\n",
    "Integrating MLflow with DialoGPT enhances conversational AI model development:\n",
    "\n",
    "- **Experiment Tracking**: Tracks configurations and metrics across experiments.\n",
    "- **Model Management**: Manages different versions and configurations of chatbot models.\n",
    "- **Reproducibility**: Ensures the reproducibility of the model's behavior.\n",
    "- **Deployment**: Simplifies deploying conversational models in production."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "env: TOKENIZERS_PARALLELISM=false\n"
     ]
    }
   ],
   "source": [
    "# Disable tokenizers warnings when constructing pipelines\n",
    "%env TOKENIZERS_PARALLELISM=false\n",
    "\n",
    "import warnings\n",
    "\n",
    "# Disable a few less-than-useful UserWarnings from setuptools and pydantic\n",
    "warnings.filterwarnings(\"ignore\", category=UserWarning)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setting Up the Conversational Pipeline\n",
    "\n",
    "We begin by setting up a conversational pipeline with DialoGPT using `transformers` and managing it with MLflow.\n",
    "\n",
    "We start by importing essential libraries. The `transformers` library from Hugging Face offers a rich collection of pre-trained models, including DialoGPT, for various NLP tasks. MLflow, a comprehensive tool for the ML lifecycle, aids in experiment tracking, reproducibility, and deployment.\n",
    "\n",
    "#### Initializing the Conversational Pipeline\n",
    "Using the `transformers.pipeline` function, we set up a conversational pipeline. We choose the \"`microsoft/DialoGPT-medium`\" model, balancing performance and resource efficiency, ideal for conversational AI. This step is pivotal for ensuring the model is ready for interaction and integration into various applications.\n",
    "\n",
    "#### Inferring the Model Signature with MLflow\n",
    "Model signature is key in defining how the model interacts with input data. To infer it, we use a sample input (\"`Hi there, chatbot!`\") and leverage `mlflow.transformers.generate_signature_output` to understand the model's input-output schema. This process ensures clarity in the model's data requirements and prediction format, crucial for seamless deployment and usage.\n",
    "\n",
    "This configuration phase sets the stage for a robust conversational AI system, leveraging the strengths of DialoGPT and MLflow for efficient and effective conversational interactions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n",
      "A decoder-only architecture is being used, but right-padding was detected! For correct generation results, please set `padding_side='left'` when initializing the tokenizer.\n"
     ]
    }
   ],
   "source": [
    "import transformers\n",
    "\n",
    "import mlflow\n",
    "\n",
    "# Define our pipeline, using the default configuration specified in the model card for DialoGPT-medium\n",
    "conversational_pipeline = transformers.pipeline(model=\"microsoft/DialoGPT-medium\")\n",
    "\n",
    "# Infer the signature by providing a representnative input and the output from the pipeline inference abstraction in the transformers flavor in MLflow\n",
    "signature = mlflow.models.infer_signature(\n",
    "    \"Hi there, chatbot!\",\n",
    "    mlflow.transformers.generate_signature_output(conversational_pipeline, \"Hi there, chatbot!\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating an experiment\n",
    "\n",
    "We create a new MLflow Experiment so that the run we're going to log our model to does not log to the default experiment and instead has its own contextually relevant entry."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<Experiment: artifact_location='file:///Users/benjamin.wilson/repos/mlflow-fork/mlflow/docs/source/llms/transformers/tutorials/conversational/mlruns/370178017237207703', creation_time=1701292102618, experiment_id='370178017237207703', last_update_time=1701292102618, lifecycle_stage='active', name='Conversational', tags={}>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# If you are running this tutorial in local mode, leave the next line commented out.\n",
    "# Otherwise, uncomment the following line and set your tracking uri to your local or remote tracking server.\n",
    "\n",
    "# mlflow.set_tracking_uri(\"http://127.0.0.1:8080\")\n",
    "\n",
    "# Set a name for the experiment that is indicative of what the runs being created within it are in regards to\n",
    "mlflow.set_experiment(\"Conversational\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Logging the Model with MLflow\n",
    "\n",
    "We'll now use MLflow to log our conversational AI model, ensuring systematic versioning, tracking, and management.\n",
    "\n",
    "#### Initiating an MLflow Run\n",
    "Our first step is to start an MLflow run with `mlflow.start_run()`. This action initiates a new tracking environment, capturing all model-related data under a unique run ID. It's a crucial step to segregate and organize different modeling experiments.\n",
    "\n",
    "#### Logging the Conversational Model\n",
    "We log our DialoGPT conversational model using `mlflow.transformers.log_model`. This specialized function efficiently logs Transformer models and requires several key parameters:\n",
    "\n",
    "- **transformers_model**: We pass our DialoGPT conversational pipeline.\n",
    "- **artifact_path**: The storage location within the MLflow run, aptly named `\"chatbot\"`.\n",
    "- **task**: Set to `\"conversational\"` to reflect the model's purpose.\n",
    "- **signature**: The inferred model signature, dictating expected inputs and outputs.\n",
    "- **input_example**: A sample prompt, like `\"A clever and witty question\"`, to demonstrate expected usage.\n",
    "\n",
    "Through this process, MLflow not only tracks our model but also organizes its metadata, facilitating future retrieval, understanding, and deployment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "with mlflow.start_run():\n",
    "    model_info = mlflow.transformers.log_model(\n",
    "        transformers_model=conversational_pipeline,\n",
    "        artifact_path=\"chatbot\",\n",
    "        task=\"conversational\",\n",
    "        signature=signature,\n",
    "        input_example=\"A clever and witty question\",\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Loading and Interacting with the Chatbot Model\n",
    "\n",
    "Next, we'll load the MLflow-logged chatbot model and interact with it to see it in action.\n",
    "\n",
    "#### Loading the Model with MLflow\n",
    "We use `mlflow.pyfunc.load_model` to load our conversational AI model. This function is a crucial aspect of MLflow's Python function flavor, offering a versatile way to interact with Python models. By specifying `model_uri=model_info.model_uri`, we precisely target the stored location of our DialoGPT model within MLflow's tracking system.\n",
    "\n",
    "#### Interacting with the Chatbot\n",
    "Once loaded, the model, referenced as `chatbot`, is ready for interaction. We demonstrate its conversational capabilities by:\n",
    "\n",
    "- **Asking Questions**: Posing a question like \"What is the best way to get to Antarctica?\" to the chatbot.\n",
    "- **Capturing Responses**: The chatbot's response, generated through the `predict` method, provides a practical example of its conversational skills. For instance, it might respond with suggestions about reaching Antarctica by boat.\n",
    "\n",
    "This demonstration highlights the practicality and convenience of deploying and using models logged with MLflow, especially in dynamic and interactive scenarios like conversational AI."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f46b4d8422cd4fac874bc5b87e85e474",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n",
      "A decoder-only architecture is being used, but right-padding was detected! For correct generation results, please set `padding_side='left'` when initializing the tokenizer.\n"
     ]
    }
   ],
   "source": [
    "# Load the model as a generic python function in order to leverage the integrated Conversational Context\n",
    "# Note that loading a conversational model with the native flavor (i.e., `mlflow.transformers.load_model()`) will not include anything apart from the\n",
    "# pipeline itself; if choosing to load in this way, you will need to manage your own Conversational Context instance to maintain state on the\n",
    "# conversation history.\n",
    "chatbot = mlflow.pyfunc.load_model(model_uri=model_info.model_uri)\n",
    "\n",
    "# Validate that the model is capable of responding to a question\n",
    "first = chatbot.predict(\"What is the best way to get to Antarctica?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Response: I think you can get there by boat.\n"
     ]
    }
   ],
   "source": [
    "print(f\"Response: {first}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Continuing the Conversation with the Chatbot\n",
    "\n",
    "We further explore the MLflow `pyfunc` implementation's conversational contextual statefulness with the DialoGPT chatbot model.\n",
    "\n",
    "#### Testing Contextual Memory\n",
    "We pose a follow-up question, \"What sort of boat should I use?\" to test the chatbot's contextual understanding. The response we get, \"A boat that can go to Antarctica,\" while straightforward, showcases the MLflow pyfunc model's ability to retain and utilize conversation history for coherent responses with `ConversationalPipeline` types of models.\n",
    "\n",
    "#### Understanding the Response Style\n",
    "The response's style – witty and slightly facetious – reflects the training data's nature, primarily conversational exchanges from Reddit. This training source significantly influences the model's tone and style, leading to responses that can be humorous and diverse.\n",
    "\n",
    "#### Implications of Training Data\n",
    "This interaction underlines the importance of the training data's source in shaping the model's responses. When deploying such models in real-world applications, it's essential to understand and consider the training data's influence on the model's conversational style and knowledge base."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n",
      "A decoder-only architecture is being used, but right-padding was detected! For correct generation results, please set `padding_side='left'` when initializing the tokenizer.\n"
     ]
    }
   ],
   "source": [
    "# Verify that the PyFunc implementation has maintained state on the conversation history by asking a vague follow-up question that requires context\n",
    "# in order to answer properly\n",
    "second = chatbot.predict(\"What sort of boat should I use?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Response: A boat that can go to Antarctica.\n"
     ]
    }
   ],
   "source": [
    "print(f\"Response: {second}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Conclusion and Key Takeaways\n",
    "\n",
    "In this tutorial, we've explored the integration of MLflow with a conversational AI model, specifically using the DialoGPT model from Microsoft. We've covered several important aspects and techniques that are crucial for anyone looking to work with advanced machine learning models in a practical, real-world setting.\n",
    "\n",
    "#### Key Takeaways\n",
    "\n",
    "1. **MLflow for Model Management**: We demonstrated how MLflow can be effectively used for managing and deploying machine learning models. The ability to log models, track experiments, and manage different versions of models is invaluable in a machine learning workflow.\n",
    "\n",
    "2. **Conversational AI**: By using the DialoGPT model, we delved into the world of conversational AI, showcasing how to set up and interact with a conversational model. This included understanding the nuances of maintaining conversational context and the impact of training data on the model's responses.\n",
    "\n",
    "3. **Practical Implementation**: Through practical examples, we showed how to log a model in MLflow, infer a model signature, and use the `pyfunc` model flavor for easy deployment and interaction. This hands-on approach is designed to provide you with the skills needed to implement these techniques in your own projects.\n",
    "\n",
    "4. **Understanding Model Responses**: We emphasized the importance of understanding the nature of the model's training data. This understanding is crucial for interpreting the model's responses and for tailoring the model to specific use cases.\n",
    "\n",
    "5. **Contextual History**: MLflow's `transformers` `pyfunc` implementation for `ConversationalPipelines` maintains a `Conversation` context without the need for managing state yourself. This enables chat bots to be created with minimal effort, since statefulness is maintained for you.\n",
    "\n",
    "### Wrapping Up\n",
    "\n",
    "As we conclude this tutorial, we hope that you have gained a deeper understanding of how to integrate MLflow with conversational AI models and the practical considerations involved in deploying these models. The skills and knowledge acquired here are not only applicable to conversational AI but also to a broader range of machine learning applications.\n",
    "\n",
    "Remember, the field of machine learning is vast and constantly evolving. Continuous learning and experimentation are key to staying updated and making the most out of these exciting technologies.\n",
    "\n",
    "Thank you for joining us in this journey through the world of MLflow and conversational AI. We encourage you to take these learnings and apply them to your own unique challenges and projects. Happy coding!\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mlflow-dev-env",
   "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.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}