{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "86b44172-7131-44a3-a825-ac6a7347b7a7",
   "metadata": {},
   "source": [
    "# TNT-LLM: Text Mining at Scale\n",
    "\n",
    "[TNT-LLM](https://arxiv.org/abs/2403.12173) by Wan, et. al describes a taxonomy generation and classification system developed by Microsoft for their Bing Copilot application.\n",
    "\n",
    "It generates a rich, interpretable taxonomy of user intents (or other categories) from raw conversation logs. This taxonomy can then be used downstream by LLMs to label logs, which in turn can be used as training data to adapt a cheap classifier (such as logistic regression classifier on embeddings) that can be deployed in your app.\n",
    "\n",
    "TNT-LLM has three main phases:\n",
    "\n",
    "1. Generate Taxonomy\n",
    "2. Label Training Data\n",
    "3. Finetune classifier + deploy\n",
    "\n",
    "When applying LangGraph in this notebook, we will focus on the first phase: taxonomy generation (blue in the diagram below). We then show how to label and fit the classifier in subsequent steps below.\n",
    "\n",
    "![TNT LLM Diagram](./img/tnt_llm.png)\n",
    "\n",
    "To generate the taxonomy, TNT-LLM proposes 5 steps:\n",
    "\n",
    "1. **Summarize** chat logs using a lower-cost LLM (batched over all logs in the sample)\n",
    "2. **Batch** the logs into random minibatches\n",
    "3. **Generate** an initial taxonomy from the first minibatch\n",
    "4. **Update** the taxonomy on each subsequent minibatch via a ritique and revise prompt\n",
    "5. **Review** the final taxonomy, scoring its quality and generating a final value using a final sample.\n",
    "\n",
    "## Prerequisites\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abd95235-4da5-4d6a-985f-78b2572ad626",
   "metadata": {},
   "outputs": [],
   "source": ["%%capture --no-stderr\n%pip install -U langgraph langchain_anthropic langsmith\n# For the embedding-based classifier use in phase 2\n%pip install -U sklearn langchain_openai"]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d98b62e4-d327-4442-8482-65529500a8a7",
   "metadata": {},
   "outputs": [],
   "source": ["import os\nfrom getpass import getpass\n\nif \"ANTHROPIC_API_KEY\" not in os.environ:\n    os.environ[\"ANTHROPIC_API_KEY\"] = getpass(\"Enter your ANTHROPIC_API_KEY: \")\n\n# (Optional) Enable tracing\nos.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\nos.environ[\"LANGCHAIN_PROJECT\"] = \"tnt-llm\"\n\nif \"LANGCHAIN_API_KEY\" not in os.environ:\n    os.environ[\"LANGCHAIN_API_KEY\"] = getpass(\"Enter your LANGCHAIN_API_KEY: \")"]
  },
  {
   "cell_type": "markdown",
   "id": "04a133e8-f94d-4ae4-8ee1-6bee09dad4fd",
   "metadata": {},
   "source": [
    "#### Graph State\n",
    "\n",
    "Since each node of a StateGraph accepts the state (and returns an updated state), we'll define that at the outset.\n",
    "\n",
    "Our flow takes in a list of documents, batches them, and then generates and refines candidate taxonomies as interpretable \"clusters\".\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "580d82b5-b60c-47a4-9c8b-e28be22ca0e3",
   "metadata": {},
   "outputs": [],
   "source": ["import logging\nimport operator\nfrom typing import Annotated, List, Optional, TypedDict\n\nlogging.basicConfig(level=logging.WARNING)\nlogger = logging.getLogger(\"tnt-llm\")\n\n\nclass Doc(TypedDict):\n    id: str\n    content: str\n    summary: Optional[str]\n    explanation: Optional[str]\n    category: Optional[str]\n\n\nclass TaxonomyGenerationState(TypedDict):\n    # The raw docs; we inject summaries within them in the first step\n    documents: List[Doc]\n    # Indices to be concise\n    minibatches: List[List[int]]\n    # Candidate Taxonomies (full trajectory)\n    clusters: Annotated[List[List[dict]], operator.add]"]
  },
  {
   "cell_type": "markdown",
   "id": "8e13d0b3-03a5-4584-98e4-b06cbb446e35",
   "metadata": {},
   "source": [
    "#### 1. Summarize Docs\n",
    "\n",
    "Chat logs can get quite long. Our taxonomy generation step needs to see large, diverse minibatches to be able to adequately capture the distribution of categories. To ensure they can all fit efficiently into the context window, we first summarize each chat log. Downstream steps will use these summaries instead of the raw doc content.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ff02c2a1-18b5-4848-96bb-27ff00978570",
   "metadata": {},
   "outputs": [],
   "source": ["import re\n\nfrom langchain import hub\nfrom langchain_anthropic import ChatAnthropic\nfrom langchain_core.output_parsers import StrOutputParser\nfrom langchain_core.runnables import RunnableConfig, RunnableLambda, RunnablePassthrough\n\nsummary_prompt = hub.pull(\"wfh/tnt-llm-summary-generation\").partial(\n    summary_length=20, explanation_length=30\n)\n\n\ndef parse_summary(xml_string: str) -> dict:\n    summary_pattern = r\"<summary>(.*?)</summary>\"\n    explanation_pattern = r\"<explanation>(.*?)</explanation>\"\n\n    summary_match = re.search(summary_pattern, xml_string, re.DOTALL)\n    explanation_match = re.search(explanation_pattern, xml_string, re.DOTALL)\n\n    summary = summary_match.group(1).strip() if summary_match else \"\"\n    explanation = explanation_match.group(1).strip() if explanation_match else \"\"\n\n    return {\"summary\": summary, \"explanation\": explanation}\n\n\nsummary_llm_chain = (\n    summary_prompt\n    | ChatAnthropic(model=\"claude-3-haiku-20240307\")\n    | StrOutputParser()\n    # Customize the tracing name for easier organization\n).with_config(run_name=\"GenerateSummary\")\nsummary_chain = summary_llm_chain | parse_summary\n\n\n# Now combine as a \"map\" operation in a map-reduce chain\n# Input: state\n# Output: state U summaries\n# Processes docs in parallel\ndef get_content(state: TaxonomyGenerationState):\n    docs = state[\"documents\"]\n    return [{\"content\": doc[\"content\"]} for doc in docs]\n\n\nmap_step = RunnablePassthrough.assign(\n    summaries=get_content\n    # This effectively creates a \"map\" operation\n    # Note you can make this more robust by handling individual errors\n    | RunnableLambda(func=summary_chain.batch, afunc=summary_chain.abatch)\n)\n\n\ndef reduce_summaries(combined: dict) -> TaxonomyGenerationState:\n    summaries = combined[\"summaries\"]\n    documents = combined[\"documents\"]\n    return {\n        \"documents\": [\n            {\n                \"id\": doc[\"id\"],\n                \"content\": doc[\"content\"],\n                \"summary\": summ_info[\"summary\"],\n                \"explanation\": summ_info[\"explanation\"],\n            }\n            for doc, summ_info in zip(documents, summaries)\n        ]\n    }\n\n\n# This is actually the node itself!\nmap_reduce_chain = map_step | reduce_summaries"]
  },
  {
   "cell_type": "markdown",
   "id": "654cbad5-3af9-4c16-9d43-cbb903957944",
   "metadata": {},
   "source": [
    "#### 2. Split into Minibatches\n",
    "\n",
    "Each minibatch contains a random sample of docs. This lets the flow identify inadequacies in the current taxonomy using new data.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "3e0139c3-b5ba-42b9-9367-33533d66eb58",
   "metadata": {},
   "outputs": [],
   "source": ["import random\n\n\ndef get_minibatches(state: TaxonomyGenerationState, config: RunnableConfig):\n    batch_size = config[\"configurable\"].get(\"batch_size\", 200)\n    original = state[\"documents\"]\n    indices = list(range(len(original)))\n    random.shuffle(indices)\n    if len(indices) < batch_size:\n        # Don't pad needlessly if we can't fill a single batch\n        return [indices]\n\n    num_full_batches = len(indices) // batch_size\n\n    batches = [\n        indices[i * batch_size : (i + 1) * batch_size] for i in range(num_full_batches)\n    ]\n\n    leftovers = len(indices) % batch_size\n    if leftovers:\n        last_batch = indices[num_full_batches * batch_size :]\n        elements_to_add = batch_size - leftovers\n        last_batch += random.sample(indices, elements_to_add)\n        batches.append(last_batch)\n\n    return {\n        \"minibatches\": batches,\n    }"]
  },
  {
   "cell_type": "markdown",
   "id": "1ca9ca18-f43a-46bd-bd28-7bec90932dd9",
   "metadata": {},
   "source": [
    "#### 3.a Taxonomy Generation Utilities\n",
    "\n",
    "This section of the graph is a generate -> update 🔄 -> review cycle. Each node shares a LOT of logic, which we have factored out into the shared functions below.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "224ed013-2963-489c-b734-315cad701d59",
   "metadata": {},
   "outputs": [],
   "source": ["from typing import Dict\n\nfrom langchain_core.runnables import Runnable\n\n\ndef parse_taxa(output_text: str) -> Dict:\n    \"\"\"Extract the taxonomy from the generated output.\"\"\"\n    cluster_matches = re.findall(\n        r\"\\s*<id>(.*?)</id>\\s*<name>(.*?)</name>\\s*<description>(.*?)</description>\\s*\",\n        output_text,\n        re.DOTALL,\n    )\n    clusters = [\n        {\"id\": id.strip(), \"name\": name.strip(), \"description\": description.strip()}\n        for id, name, description in cluster_matches\n    ]\n    # We don't parse the explanation since it isn't used downstream\n    return {\"clusters\": clusters}\n\n\ndef format_docs(docs: List[Doc]) -> str:\n    xml_table = \"<conversations>\\n\"\n    for doc in docs:\n        xml_table += f'<conv_summ id={doc[\"id\"]}>{doc[\"summary\"]}</conv_summ>\\n'\n    xml_table += \"</conversations>\"\n    return xml_table\n\n\ndef format_taxonomy(clusters):\n    xml = \"<cluster_table>\\n\"\n    for label in clusters:\n        xml += \"  <cluster>\\n\"\n        xml += f'    <id>{label[\"id\"]}</id>\\n'\n        xml += f'    <name>{label[\"name\"]}</name>\\n'\n        xml += f'    <description>{label[\"description\"]}</description>\\n'\n        xml += \"  </cluster>\\n\"\n    xml += \"</cluster_table>\"\n    return xml\n\n\ndef invoke_taxonomy_chain(\n    chain: Runnable,\n    state: TaxonomyGenerationState,\n    config: RunnableConfig,\n    mb_indices: List[int],\n) -> TaxonomyGenerationState:\n    configurable = config[\"configurable\"]\n    docs = state[\"documents\"]\n    minibatch = [docs[idx] for idx in mb_indices]\n    data_table_xml = format_docs(minibatch)\n\n    previous_taxonomy = state[\"clusters\"][-1] if state[\"clusters\"] else []\n    cluster_table_xml = format_taxonomy(previous_taxonomy)\n\n    updated_taxonomy = chain.invoke(\n        {\n            \"data_xml\": data_table_xml,\n            \"use_case\": configurable[\"use_case\"],\n            \"cluster_table_xml\": cluster_table_xml,\n            \"suggestion_length\": configurable.get(\"suggestion_length\", 30),\n            \"cluster_name_length\": configurable.get(\"cluster_name_length\", 10),\n            \"cluster_description_length\": configurable.get(\n                \"cluster_description_length\", 30\n            ),\n            \"explanation_length\": configurable.get(\"explanation_length\", 20),\n            \"max_num_clusters\": configurable.get(\"max_num_clusters\", 25),\n        }\n    )\n\n    return {\n        \"clusters\": [updated_taxonomy[\"clusters\"]],\n    }"]
  },
  {
   "cell_type": "markdown",
   "id": "2e2a2723-d350-4871-83e8-88f081ab4c8b",
   "metadata": {},
   "source": [
    "#### 3. Generate initial taxonomy\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "553dff30-ce53-47d8-ab3c-d2f437b7d5f4",
   "metadata": {},
   "outputs": [],
   "source": ["# We will share an LLM for each step of the generate -> update -> review cycle\n# You may want to consider using Opus or another more powerful model for this\ntaxonomy_generation_llm = ChatAnthropic(\n    model=\"claude-3-haiku-20240307\", max_tokens_to_sample=2000\n)\n\n\n## Initial generation\ntaxonomy_generation_prompt = hub.pull(\"wfh/tnt-llm-taxonomy-generation\").partial(\n    use_case=\"Generate the taxonomy that can be used to label the user intent in the conversation.\",\n)\n\ntaxa_gen_llm_chain = (\n    taxonomy_generation_prompt | taxonomy_generation_llm | StrOutputParser()\n).with_config(run_name=\"GenerateTaxonomy\")\n\n\ngenerate_taxonomy_chain = taxa_gen_llm_chain | parse_taxa\n\n\ndef generate_taxonomy(\n    state: TaxonomyGenerationState, config: RunnableConfig\n) -> TaxonomyGenerationState:\n    return invoke_taxonomy_chain(\n        generate_taxonomy_chain, state, config, state[\"minibatches\"][0]\n    )"]
  },
  {
   "cell_type": "markdown",
   "id": "8caefddd-c4c6-4318-972a-39c384183d79",
   "metadata": {},
   "source": [
    "#### 4. Update Taxonomy\n",
    "\n",
    "This is a \"critique -> revise\" step that is repeated N times.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "b8739b5b-ba8a-4c40-bd25-a3b06a19949d",
   "metadata": {},
   "outputs": [],
   "source": ["taxonomy_update_prompt = hub.pull(\"wfh/tnt-llm-taxonomy-update\")\n\ntaxa_update_llm_chain = (\n    taxonomy_update_prompt | taxonomy_generation_llm | StrOutputParser()\n).with_config(run_name=\"UpdateTaxonomy\")\n\n\nupdate_taxonomy_chain = taxa_update_llm_chain | parse_taxa\n\n\ndef update_taxonomy(\n    state: TaxonomyGenerationState, config: RunnableConfig\n) -> TaxonomyGenerationState:\n    which_mb = len(state[\"clusters\"]) % len(state[\"minibatches\"])\n    return invoke_taxonomy_chain(\n        update_taxonomy_chain, state, config, state[\"minibatches\"][which_mb]\n    )"]
  },
  {
   "cell_type": "markdown",
   "id": "35f07f42-3025-446c-a6b0-2e45ea54f269",
   "metadata": {},
   "source": [
    "#### 5. Review Taxonomy\n",
    "\n",
    "This runs once we've processed all the minibatches.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "0039cf1c-54d5-4e9e-8dd6-a5cebfaec92d",
   "metadata": {},
   "outputs": [],
   "source": ["taxonomy_review_prompt = hub.pull(\"wfh/tnt-llm-taxonomy-review\")\n\ntaxa_review_llm_chain = (\n    taxonomy_review_prompt | taxonomy_generation_llm | StrOutputParser()\n).with_config(run_name=\"ReviewTaxonomy\")\n\n\nreview_taxonomy_chain = taxa_review_llm_chain | parse_taxa\n\n\ndef review_taxonomy(\n    state: TaxonomyGenerationState, config: RunnableConfig\n) -> TaxonomyGenerationState:\n    batch_size = config[\"configurable\"].get(\"batch_size\", 200)\n    original = state[\"documents\"]\n    indices = list(range(len(original)))\n    random.shuffle(indices)\n    return invoke_taxonomy_chain(\n        review_taxonomy_chain, state, config, indices[:batch_size]\n    )"]
  },
  {
   "cell_type": "markdown",
   "id": "ae1d8103-3ecb-458c-8269-1f81c1c6296b",
   "metadata": {},
   "source": [
    "## Define the Graph\n",
    "\n",
    "With all the functionality defined, we can define the graph!\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "f1f97ea4-53e5-4f55-8d73-b5b2234a47d9",
   "metadata": {},
   "outputs": [],
   "source": ["from langgraph.graph import StateGraph, START, END\n\ngraph = StateGraph(TaxonomyGenerationState)\ngraph.add_node(\"summarize\", map_reduce_chain)\ngraph.add_node(\"get_minibatches\", get_minibatches)\ngraph.add_node(\"generate_taxonomy\", generate_taxonomy)\ngraph.add_node(\"update_taxonomy\", update_taxonomy)\ngraph.add_node(\"review_taxonomy\", review_taxonomy)\n\ngraph.add_edge(\"summarize\", \"get_minibatches\")\ngraph.add_edge(\"get_minibatches\", \"generate_taxonomy\")\ngraph.add_edge(\"generate_taxonomy\", \"update_taxonomy\")\n\n\ndef should_review(state: TaxonomyGenerationState) -> str:\n    num_minibatches = len(state[\"minibatches\"])\n    num_revisions = len(state[\"clusters\"])\n    if num_revisions < num_minibatches:\n        return \"update_taxonomy\"\n    return \"review_taxonomy\"\n\n\ngraph.add_conditional_edges(\n    \"update_taxonomy\",\n    should_review,\n    # Optional (but required for the diagram to be drawn correctly below)\n    {\"update_taxonomy\": \"update_taxonomy\", \"review_taxonomy\": \"review_taxonomy\"},\n)\ngraph.add_edge(\"review_taxonomy\", END)\n\ngraph.add_edge(START, \"summarize\")\napp = graph.compile()"]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "cc4fcd31-a380-4eac-872a-42edd93736c6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 84124 characters]",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": ["from IPython.display import Image\n\nImage(app.get_graph().draw_png())"]
  },
  {
   "cell_type": "markdown",
   "id": "e8bdebf5-f315-4e96-80a2-5dde66b327c1",
   "metadata": {},
   "source": [
    "## Usage\n",
    "\n",
    "The docs can contain **any** content, but we've found it works really well on chat bot logs, such as those captured by [LangSmith](https://smith.langchain.com).\n",
    "\n",
    "We will use that as an example below. Update the `project_name` to your own LangSmith project.\n",
    "\n",
    "You will likely have to customize the `run_to_doc` function below, since your expected keys may differ from those of this notebook's author.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 193,
   "id": "bcc65649-157f-4848-9ef0-8a9932a98d85",
   "metadata": {},
   "outputs": [],
   "source": ["from datetime import datetime, timedelta\n\nfrom langsmith import Client\n\nproject_name = \"YOUR PROJECT NAME\"  # Update to your own project\nclient = Client()\n\npast_week = datetime.now() - timedelta(days=7)\nruns = list(\n    client.list_runs(\n        project_name=project_name,\n        filter=\"eq(is_root, true)\",\n        start_time=past_week,\n        # We only need to return the inputs + outputs\n        select=[\"inputs\", \"outputs\"],\n    )\n)\n\n\n# Convert the langsmith traces to our graph's Doc object.\ndef run_to_doc(run) -> Doc:\n    turns = []\n    idx = 0\n    for turn in run.inputs.get(\"chat_history\") or []:\n        key, value = next(iter(turn.items()))\n        turns.append(f\"<{key} idx={idx}>\\n{value}\\n</{key}>\")\n        idx += 1\n    turns.append(\n        f\"\"\"\n<human idx={idx}>\n{run.inputs['question']}\n</human>\"\"\"\n    )\n    if run.outputs and run.outputs[\"output\"]:\n        turns.append(\n            f\"\"\"<ai idx={idx+1}>\n{run.outputs['output']}\n</ai>\"\"\"\n        )\n    return {\n        \"id\": str(run.id),\n        \"content\": (\"\\n\".join(turns)),\n    }"]
  },
  {
   "cell_type": "markdown",
   "id": "245ed4e4-c08d-4e47-90d5-b437b0532954",
   "metadata": {},
   "source": [
    "#### Invoke\n",
    "\n",
    "Now convert the runs to docs and kick off your graph flow. This will take some time! The summary step takes the longest. If you want to speed things up, you could try splitting the load across model providers.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "900906b5-9264-46a8-ba83-46307f8c25d0",
   "metadata": {},
   "outputs": [],
   "source": ["from langchain.cache import InMemoryCache\nfrom langchain.globals import set_llm_cache\n\n# Optional. If you are running into errors or rate limits and want to avoid repeated computation,\n# you can set this while debugging\n\nset_llm_cache(InMemoryCache())"]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2340177-f40c-407a-8e3e-cb06c2ef09ce",
   "metadata": {},
   "outputs": [],
   "source": ["# We will randomly sample down to 1K docs to speed things up\ndocs = [run_to_doc(run) for run in runs if run.inputs]\ndocs = random.sample(docs, min(len(docs), 1000))\nuse_case = (\n    \"Generate the taxonomy that can be used both to label the user intent\"\n    \" as well as to identify any required documentation (references, how-tos, etc.)\"\n    \" that would benefit the user.\"\n)\n\nstream = app.stream(\n    {\"documents\": docs},\n    {\n        \"configurable\": {\n            \"use_case\": use_case,\n            # Optional:\n            \"batch_size\": 400,\n            \"suggestion_length\": 30,\n            \"cluster_name_length\": 10,\n            \"cluster_description_length\": 30,\n            \"explanation_length\": 20,\n            \"max_num_clusters\": 25,\n        },\n        # We batch summarize the docs. To avoid getting errors, we will limit the\n        # degree of parallelism to permit.\n        \"max_concurrency\": 2,\n    },\n)\n\nfor step in stream:\n    node, state = next(iter(step.items()))\n    print(node, str(state)[:20] + \" ...\")"]
  },
  {
   "cell_type": "markdown",
   "id": "46be818a-adda-4ecc-a798-7074eb1ec400",
   "metadata": {},
   "source": [
    "## Final Result\n",
    "\n",
    "Below, render the final result as markdown:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 202,
   "id": "38085515-6cda-452d-9f44-de440058cb3f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "## Final Taxonomy\n",
       "\n",
       "| ID | Name | Description |\n",
       "|----|------|-------------|\n",
       "| 1 | Troubleshooting Network Connectivity Issues | Resolving problems with DNS, network connections, and GitHub extension activation. |\n",
       "| 2 | Extracting and Analyzing Data | Retrieving and processing data from various sources like text files, databases, and APIs. |\n",
       "| 3 | Providing Healthcare Insights | Generating medical diagnosis, symptom checking, drug information, and skin condition analysis. |\n",
       "| 4 | Configuring and Optimizing Models | Adjusting model parameters and hyperparameters to improve performance for a given task. |\n",
       "| 5 | Generating Creative Poetry | Creating poems using language models and AI-powered tools. |\n",
       "| 6 | Interacting with Databases | Querying databases, extracting data, and managing errors during data processing. |\n",
       "| 7 | Querying Vector Databases | Interacting with vector databases like Milvus to store and retrieve high-dimensional data. |\n",
       "| 8 | Generating Synthetic Data | Creating synthetic data using language models and machine learning techniques. |\n",
       "| 9 | Integrating Tools and Workflows | Incorporating various tools and libraries into a cohesive workflow for different tasks. |\n",
       "| 10 | Improving Information Retrieval | Storing and querying multiple vectors per document for better semantic understanding. |\n",
       "| 11 | Processing Documents and Extracting Text | Parsing and extracting text from various document formats like PDF, DOCX, and HTML. |\n",
       "| 12 | Building Local Knowledge Bases | Creating knowledge bases from text files, handling text splitting, embeddings, and storage. |\n",
       "| 13 | Optimizing Conversational Retrieval | Troubleshooting and improving the performance of the ConversationalRetrievalChain in LangChain. |\n",
       "| 14 | Connecting Databases and Using Agents | Connecting to databases, using agents, and understanding the differences between agent types. |\n",
       "| 15 | Introspecting LangChain Tools | Accessing and retrieving details about the functions and source code of LangChain tools. |\n",
       "| 16 | Generating Styled Answers with Retrieval Augmentation | Creating a QA system that generates well-cited answers in a specific style. |\n",
       "| 17 | Using ZERO_SHOT_REACT_DESCRIPTION Agents | Applying the ZERO_SHOT_REACT_DESCRIPTION agent type in LangChain for chat models. |\n",
       "| 18 | Automating Microlearning Course Creation | Generating microlearning courses based on input parameters like topic, volume, and learning style. |\n",
       "| 19 | Integrating with Chroma Vector Store | Storing and retrieving data in the Chroma vector database, including handling document embeddings. |\n",
       "| 20 | Managing LangChain Callback Tokens | Understanding and utilizing the callback token feature in the LCEL chain. |\n",
       "| 21 | Troubleshooting FastAPI Deployments | Resolving issues with deploying a React app with a FastAPI backend. |\n",
       "| 22 | Analyzing Data with LangChain Agents | Using LangChain agents to interact with Pandas and Spark DataFrames for data exploration. |\n",
       "| 23 | Implementing the OpenAI Chat API | Implementing the OpenAI chat completion API and understanding the required inputs and outputs. |\n",
       "| 24 | Comparing LangChain and LLMIndex | Evaluating the differences between LangChain and LLMIndex, including their UI support for Markdown. |\n",
       "| 25 | Suppressing Tools in AgentExecutor | Temporarily disabling tools in an AgentExecutor for a fixed number of invocations. |\n"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "execution_count": 202,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": ["from IPython.display import Markdown\n\n\ndef format_taxonomy_md(clusters):\n    md = \"## Final Taxonomy\\n\\n\"\n    md += \"| ID | Name | Description |\\n\"\n    md += \"|----|------|-------------|\\n\"\n\n    # Fill the table with cluster data\n    for label in clusters:\n        id = label[\"id\"]\n        name = label[\"name\"].replace(\n            \"|\", \"\\\\|\"\n        )  # Escape any pipe characters within the content\n        description = label[\"description\"].replace(\n            \"|\", \"\\\\|\"\n        )  # Escape any pipe characters\n        md += f\"| {id} | {name} | {description} |\\n\"\n\n    return md\n\n\nMarkdown(format_taxonomy_md(step[\"__end__\"][\"clusters\"][-1]))"]
  },
  {
   "cell_type": "markdown",
   "id": "3cf106db-6cff-4847-bff9-e25f00b9c015",
   "metadata": {},
   "source": [
    "## Phase 2: Labeling\n",
    "\n",
    "Now that we have our taxonomy, it's time to label a subset of our data to train a classifier.\n",
    "\n",
    "Input classification can be useful for anything from in-line prompt optimization (tailor the prompt for each classified intent), to system improvements (identifying categories for which the system doesn't produce good responses) to product analytics (understand which intent categories could be improved to drive profits).\n",
    "\n",
    "The problem is that LLM-based tagging can be expensive.\n",
    "\n",
    "Embeddings can be ~100x cheaper to compute, and a simple logistic regression classifier on top of that would add negligible cost.\n",
    "\n",
    "Let's tag and train a classifier!\n",
    "\n",
    "#### Label Training Data\n",
    "\n",
    "Use an LLM to label the data in a fully-automated fashion. For better accuracy, you can sample a portion of the results to label by hand as well to verify the quality.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "id": "8aa8a6f5-f53a-41e5-b09d-c6e8476e5471",
   "metadata": {},
   "outputs": [],
   "source": ["labeling_prompt = hub.pull(\"wfh/tnt-llm-classify\")\n\nlabeling_llm = ChatAnthropic(model=\"claude-3-haiku-20240307\", max_tokens_to_sample=2000)\nlabeling_llm_chain = (labeling_prompt | labeling_llm | StrOutputParser()).with_config(\n    run_name=\"ClassifyDocs\"\n)\n\n\ndef parse_labels(output_text: str) -> Dict:\n    \"\"\"Parse the generated labels from the predictions.\"\"\"\n    category_matches = re.findall(\n        r\"\\s*<category>(.*?)</category>.*\",\n        output_text,\n        re.DOTALL,\n    )\n    categories = [{\"category\": category.strip()} for category in category_matches]\n    if len(categories) > 1:\n        logger.warning(f\"Multiple selected categories: {categories}\")\n    label = categories[0]\n    stripped = re.sub(r\"^\\d+\\.\\s*\", \"\", label[\"category\"]).strip()\n    return {\"category\": stripped}\n\n\nlabeling_chain = labeling_llm_chain | parse_labels"]
  },
  {
   "cell_type": "code",
   "execution_count": 148,
   "id": "59c06eea-ecbf-43af-a292-71816ccd92b8",
   "metadata": {},
   "outputs": [],
   "source": ["final_taxonomy = step[\"__end__\"][\"clusters\"][-1]\nxml_taxonomy = format_taxonomy(final_taxonomy)\nresults = labeling_chain.batch(\n    [\n        {\n            \"content\": doc[\"content\"],\n            \"taxonomy\": xml_taxonomy,\n        }\n        for doc in docs\n    ],\n    {\"max_concurrency\": 5},\n    return_exceptions=True,\n)\n# Update the docs to include the categories\nupdated_docs = [{**doc, **category} for doc, category in zip(docs, results)]"]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ef9be82-278e-4501-8af9-70409ce15cc2",
   "metadata": {},
   "outputs": [],
   "source": ["if \"OPENAI_API_KEY\" not in os.environ:\n    os.environ[\"OPENAI_API_KEY\"] = getpass(\"Enter your OPENAI_API_KEY: \")"]
  },
  {
   "cell_type": "code",
   "execution_count": 149,
   "id": "c21f787e-2dcb-49c2-9cc1-5284a1732fbc",
   "metadata": {},
   "outputs": [],
   "source": ["from langchain_openai import OpenAIEmbeddings\n\n# Consider using other embedding models here too!\nencoder = OpenAIEmbeddings(model=\"text-embedding-3-large\")\nvectors = encoder.embed_documents([doc[\"content\"] for doc in docs])\nembedded_docs = [{**doc, \"embedding\": v} for doc, v in zip(updated_docs, vectors)]"]
  },
  {
   "cell_type": "markdown",
   "id": "72284867-727d-467a-891a-7c7881967675",
   "metadata": {},
   "source": [
    "#### Train Classifier\n",
    "\n",
    "Now that we've extracted the features from the text, we can generate the classifier on them.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 196,
   "id": "eb19d42b-97c9-466b-b7b3-c17cc4db5180",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Accuracy: 0.515\n",
      "Test Accuracy: 0.330\n",
      "Train F1 Score: 0.493\n",
      "Test F1 Score: 0.335\n"
     ]
    }
   ],
   "source": ["import numpy as np\nfrom sklearn.linear_model import LogisticRegression\nfrom sklearn.metrics import accuracy_score, f1_score\nfrom sklearn.model_selection import train_test_split\nfrom sklearn.utils import class_weight\n\n# Create a dictionary mapping category names to their indices in the taxonomy\ncategory_to_index = {d[\"name\"]: i for i, d in enumerate(final_taxonomy)}\ncategory_to_index[\"Other\"] = len(category_to_index)\n# Convert category strings to numeric labels\nlabels = [\n    category_to_index.get(d[\"category\"], category_to_index[\"Other\"])\n    for d in embedded_docs\n]\n\nlabel_vectors = [d[\"embedding\"] for d in embedded_docs]\n\nX_train, X_test, y_train, y_test = train_test_split(\n    label_vectors, labels, test_size=0.2, random_state=42\n)\n\n# Calculate class weights\nclass_weights = class_weight.compute_class_weight(\n    class_weight=\"balanced\", classes=np.unique(y_train), y=y_train\n)\nclass_weight_dict = dict(enumerate(class_weights))\n\n# Weight the classes to partially handle imbalanced data\nmodel = LogisticRegression(class_weight=class_weight_dict)\nmodel.fit(X_train, y_train)\n\ntrain_preds = model.predict(X_train)\ntest_preds = model.predict(X_test)\n\ntrain_acc = accuracy_score(y_train, train_preds)\ntest_acc = accuracy_score(y_test, test_preds)\ntrain_f1 = f1_score(y_train, train_preds, average=\"weighted\")\ntest_f1 = f1_score(y_test, test_preds, average=\"weighted\")\n\nprint(f\"Train Accuracy: {train_acc:.3f}\")\nprint(f\"Test Accuracy: {test_acc:.3f}\")\nprint(f\"Train F1 Score: {train_f1:.3f}\")\nprint(f\"Test F1 Score: {test_f1:.3f}\")"]
  },
  {
   "cell_type": "markdown",
   "id": "32e7e994-a3c3-47f3-a860-32b847dfefee",
   "metadata": {},
   "source": [
    "## Phase 3: Deploy\n",
    "\n",
    "Now that you have your classifier, you can easily deploy it and apply to future runs! All you need is to embed the input and apply your LogisticRegression classifier. Let's try it. We will use python's [joblib](https://joblib.readthedocs.io/en/stable/) library to serialize our sklearn classifier. Below is an example:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 197,
   "id": "c27cbb6b-4d0f-476a-bef3-31ed307ce45f",
   "metadata": {},
   "outputs": [],
   "source": ["from joblib import dump as jl_dump\n\ncategories = list(category_to_index)\n\n# Save the model and categories to a file\nwith open(\"model.joblib\", \"wb\") as file:\n    jl_dump((model, categories), file)"]
  },
  {
   "cell_type": "markdown",
   "id": "443f5f92-7f3c-4ce8-9104-b2b32ace0c42",
   "metadata": {},
   "source": [
    "#### To deploy\n",
    "\n",
    "When deploying, you can load the classifier and initialize your embeddings encoder. They fit together easily using LCEL:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 198,
   "id": "28f0b88a-b308-4208-b482-6c157357dfc6",
   "metadata": {},
   "outputs": [],
   "source": ["from joblib import load as jl_load\nfrom langchain_openai import OpenAIEmbeddings\n\nloaded_model, loaded_categories = jl_load(\"model.joblib\")\nencoder = OpenAIEmbeddings(model=\"text-embedding-3-large\")\n\n\ndef get_category_name(predictions):\n    return [loaded_categories[pred] for pred in predictions]\n\n\nclassifier = (\n    RunnableLambda(encoder.embed_documents, encoder.aembed_documents)\n    | loaded_model.predict\n    | get_category_name\n)"]
  },
  {
   "cell_type": "markdown",
   "id": "3114657e-b558-4a72-91af-68ae3fccacd7",
   "metadata": {},
   "source": [
    "#### Example:\n",
    "\n",
    "Assuming you've had some more data come in, you can fetch it and apply it below\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 194,
   "id": "6cdb9d8a-2aa1-4f48-8b23-f311fdf36416",
   "metadata": {},
   "outputs": [],
   "source": ["client = Client()\n\npast_5_min = datetime.now() - timedelta(minutes=5)\nruns = list(\n    client.list_runs(\n        project_name=project_name,\n        filter=\"eq(is_root, true)\",\n        start_time=past_5_min,\n        # We only need to return the inputs + outputs\n        select=[\"inputs\", \"outputs\"],\n        limit=100,\n    )\n)\ndocs = [run_to_doc(r) for r in runs]"]
  },
  {
   "cell_type": "code",
   "execution_count": 199,
   "id": "b154439b-902e-42f3-afbc-89edb0cfea81",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:httpx:HTTP Request: POST https://api.openai.com/v1/embeddings \"HTTP/1.1 200 OK\"\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['Interacting with Databases', 'Optimizing Conversational Retrieval']\n"
     ]
    }
   ],
   "source": ["classes = classifier.invoke([doc[\"content\"] for doc in docs])\nprint(classes[:2])"]
  },
  {
   "cell_type": "markdown",
   "id": "fe415d54-fa94-4c62-b68b-d4437c446d42",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "Congrats on implementing TNT-LLM! While most folks use clustering-based approaches like LDA, k-means, etc. it can often be hard to really interpret what each cluster represents. TNT-LLM generates human-interpretable labels you can use downstream to monitor and improve your application.\n",
    "\n",
    "The technique also lends itself to hierarchical sub-categorizing: once you have the above taxonomy, use it to label your data, then on each sub-category, generate a new taxonomy using a similar technique to the one described above!\n"
   ]
  }
 ],
 "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.11.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
