{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "70996d8a",
   "metadata": {},
   "source": [
    "# WatsonxLLM\n",
    "\n",
    "[WatsonxLLM](https://ibm.github.io/watson-machine-learning-sdk/fm_extensions.html) is wrapper for IBM [watsonx.ai](https://www.ibm.com/products/watsonx-ai) foundation models.\n",
    "This example shows how to communicate with watsonx.ai models using LangChain."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea35b2b7",
   "metadata": {},
   "source": [
    "Install the package [`ibm_watson_machine_learning`](https://ibm.github.io/watson-machine-learning-sdk/install.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f1fff4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install ibm_watson_machine_learning"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f406e092",
   "metadata": {},
   "source": [
    "This cell defines the WML credentials required to work with watsonx Foundation Model inferencing.\n",
    "\n",
    "**Action:** Provide the IBM Cloud user API key. For details, see\n",
    "[documentation](https://cloud.ibm.com/docs/account?topic=account-userapikey&interface=ui)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "11d572a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from getpass import getpass\n",
    "\n",
    "watsonx_api_key = getpass()\n",
    "os.environ[\"WATSONX_APIKEY\"] = watsonx_api_key"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e36acbef",
   "metadata": {},
   "source": [
    "## Load the model\n",
    "You might need to adjust model `parameters` for different models or tasks, to do so please refer to [documentation](https://ibm.github.io/watson-machine-learning-sdk/model.html#metanames.GenTextParamsMetaNames)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "407cd500",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ibm_watson_machine_learning.metanames import GenTextParamsMetaNames as GenParams\n",
    "\n",
    "parameters = {\n",
    "    GenParams.DECODING_METHOD: \"sample\",\n",
    "    GenParams.MAX_NEW_TOKENS: 100,\n",
    "    GenParams.MIN_NEW_TOKENS: 1,\n",
    "    GenParams.TEMPERATURE: 0.5,\n",
    "    GenParams.TOP_K: 50,\n",
    "    GenParams.TOP_P: 1,\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b586538",
   "metadata": {},
   "source": [
    "Initialize the `WatsonxLLM` class with previous set params."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "359898de",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.llms import WatsonxLLM\n",
    "\n",
    "watsonx_llm = WatsonxLLM(\n",
    "    model_id=\"google/flan-ul2\",\n",
    "    url=\"https://us-south.ml.cloud.ibm.com\",\n",
    "    project_id=\"***\",\n",
    "    params=parameters,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2202f4e0",
   "metadata": {},
   "source": [
    "Alternatively you can use Cloud Pak for Data credentials. For details, see [documentation](https://ibm.github.io/watson-machine-learning-sdk/setup_cpd.html).\n",
    "```\n",
    "watsonx_llm = WatsonxLLM(\n",
    "    model_id='google/flan-ul2',\n",
    "    url=\"***\",\n",
    "    username=\"***\",\n",
    "    password=\"***\",\n",
    "    instance_id=\"openshift\",\n",
    "    version=\"4.8\",\n",
    "    project_id='***',\n",
    "    params=parameters\n",
    ")\n",
    "```        "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c25ecbd1",
   "metadata": {},
   "source": [
    "## Create Chain\n",
    "Create `PromptTemplate` objects which will be responsible for creating a random question."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c7d80c05",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.prompts import PromptTemplate\n",
    "\n",
    "template = \"Generate a random question about {topic}: Question: \"\n",
    "prompt = PromptTemplate.from_template(template)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79056d8e",
   "metadata": {},
   "source": [
    "Provide a topic and run the `LLMChain`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "dc076c56",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'How many breeds of dog are there?'"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain.chains import LLMChain\n",
    "\n",
    "llm_chain = LLMChain(prompt=prompt, llm=watsonx_llm)\n",
    "llm_chain.run(\"dog\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f571001d",
   "metadata": {},
   "source": [
    "## Calling the Model Directly\n",
    "To obtain completions, you can can the model directly using string prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "beea2b5b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'dog'"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Calling a single prompt\n",
    "\n",
    "watsonx_llm(\"Who is man's best friend?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "8ab1a25a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "LLMResult(generations=[[Generation(text='greyhounds', generation_info={'generated_token_count': 4, 'input_token_count': 8, 'finish_reason': 'eos_token'})], [Generation(text='The Basenji is a dog breed from South Africa.', generation_info={'generated_token_count': 13, 'input_token_count': 7, 'finish_reason': 'eos_token'})]], llm_output={'model_id': 'google/flan-ul2'}, run=[RunInfo(run_id=UUID('03c73a42-db68-428e-ab8d-8ae10abc84fc')), RunInfo(run_id=UUID('c289f67a-87d6-4c8b-a8b7-0b5012c94ca8'))])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Calling multiple prompts\n",
    "\n",
    "watsonx_llm.generate(\n",
    "    [\n",
    "        \"The fastest dog in the world?\",\n",
    "        \"Describe your chosen dog breed\",\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2c9da33",
   "metadata": {},
   "source": [
    "## Streaming the Model output \n",
    "\n",
    "You can stream the model output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "3f63166a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The golden retriever is my favorite dog because it is very friendly and good with children."
     ]
    }
   ],
   "source": [
    "for chunk in watsonx_llm.stream(\n",
    "    \"Describe your favorite breed of dog and why it is your favorite.\"\n",
    "):\n",
    "    print(chunk, end=\"\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
