[Docs] Update notebooks to use START (#902)

This commit is contained in:
William FH
2024-07-01 21:36:34 -07:00
committed by GitHub
parent 267f5e5234
commit 727e63c01e
67 changed files with 1059 additions and 20258 deletions
File diff suppressed because one or more lines are too long
+24 -587
View File
@@ -38,12 +38,7 @@
"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"
]
"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",
@@ -51,20 +46,7 @@
"id": "d98b62e4-d327-4442-8482-65529500a8a7",
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"from getpass import getpass\n",
"\n",
"if \"ANTHROPIC_API_KEY\" not in os.environ:\n",
" os.environ[\"ANTHROPIC_API_KEY\"] = getpass(\"Enter your ANTHROPIC_API_KEY: \")\n",
"\n",
"# (Optional) Enable tracing\n",
"os.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\n",
"os.environ[\"LANGCHAIN_PROJECT\"] = \"tnt-llm\"\n",
"\n",
"if \"LANGCHAIN_API_KEY\" not in os.environ:\n",
" os.environ[\"LANGCHAIN_API_KEY\"] = getpass(\"Enter your LANGCHAIN_API_KEY: \")"
]
"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",
@@ -84,31 +66,7 @@
"id": "580d82b5-b60c-47a4-9c8b-e28be22ca0e3",
"metadata": {},
"outputs": [],
"source": [
"import logging\n",
"import operator\n",
"from typing import Annotated, List, Optional, TypedDict\n",
"\n",
"logging.basicConfig(level=logging.WARNING)\n",
"logger = logging.getLogger(\"tnt-llm\")\n",
"\n",
"\n",
"class Doc(TypedDict):\n",
" id: str\n",
" content: str\n",
" summary: Optional[str]\n",
" explanation: Optional[str]\n",
" category: Optional[str]\n",
"\n",
"\n",
"class 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]"
]
"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",
@@ -126,77 +84,7 @@
"id": "ff02c2a1-18b5-4848-96bb-27ff00978570",
"metadata": {},
"outputs": [],
"source": [
"import re\n",
"\n",
"from langchain import hub\n",
"from langchain_anthropic import ChatAnthropic\n",
"from langchain_core.output_parsers import StrOutputParser\n",
"from langchain_core.runnables import RunnableConfig, RunnableLambda, RunnablePassthrough\n",
"\n",
"summary_prompt = hub.pull(\"wfh/tnt-llm-summary-generation\").partial(\n",
" summary_length=20, explanation_length=30\n",
")\n",
"\n",
"\n",
"def 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",
"\n",
"summary_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\")\n",
"summary_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\n",
"def get_content(state: TaxonomyGenerationState):\n",
" docs = state[\"documents\"]\n",
" return [{\"content\": doc[\"content\"]} for doc in docs]\n",
"\n",
"\n",
"map_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",
"\n",
"def 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!\n",
"map_reduce_chain = map_step | reduce_summaries"
]
"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",
@@ -214,36 +102,7 @@
"id": "3e0139c3-b5ba-42b9-9367-33533d66eb58",
"metadata": {},
"outputs": [],
"source": [
"import random\n",
"\n",
"\n",
"def 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",
" }"
]
"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",
@@ -261,80 +120,7 @@
"id": "224ed013-2963-489c-b734-315cad701d59",
"metadata": {},
"outputs": [],
"source": [
"from typing import Dict\n",
"\n",
"from langchain_core.runnables import Runnable\n",
"\n",
"\n",
"def 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",
"\n",
"def 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",
"\n",
"def 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",
"\n",
"def 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",
" }"
]
"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",
@@ -350,34 +136,7 @@
"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\n",
"taxonomy_generation_llm = ChatAnthropic(\n",
" model=\"claude-3-haiku-20240307\", max_tokens_to_sample=2000\n",
")\n",
"\n",
"\n",
"## Initial generation\n",
"taxonomy_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",
"\n",
"taxa_gen_llm_chain = (\n",
" taxonomy_generation_prompt | taxonomy_generation_llm | StrOutputParser()\n",
").with_config(run_name=\"GenerateTaxonomy\")\n",
"\n",
"\n",
"generate_taxonomy_chain = taxa_gen_llm_chain | parse_taxa\n",
"\n",
"\n",
"def generate_taxonomy(\n",
" state: TaxonomyGenerationState, config: RunnableConfig\n",
") -> TaxonomyGenerationState:\n",
" return invoke_taxonomy_chain(\n",
" generate_taxonomy_chain, state, config, state[\"minibatches\"][0]\n",
" )"
]
"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",
@@ -395,25 +154,7 @@
"id": "b8739b5b-ba8a-4c40-bd25-a3b06a19949d",
"metadata": {},
"outputs": [],
"source": [
"taxonomy_update_prompt = hub.pull(\"wfh/tnt-llm-taxonomy-update\")\n",
"\n",
"taxa_update_llm_chain = (\n",
" taxonomy_update_prompt | taxonomy_generation_llm | StrOutputParser()\n",
").with_config(run_name=\"UpdateTaxonomy\")\n",
"\n",
"\n",
"update_taxonomy_chain = taxa_update_llm_chain | parse_taxa\n",
"\n",
"\n",
"def 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",
" )"
]
"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",
@@ -431,28 +172,7 @@
"id": "0039cf1c-54d5-4e9e-8dd6-a5cebfaec92d",
"metadata": {},
"outputs": [],
"source": [
"taxonomy_review_prompt = hub.pull(\"wfh/tnt-llm-taxonomy-review\")\n",
"\n",
"taxa_review_llm_chain = (\n",
" taxonomy_review_prompt | taxonomy_generation_llm | StrOutputParser()\n",
").with_config(run_name=\"ReviewTaxonomy\")\n",
"\n",
"\n",
"review_taxonomy_chain = taxa_review_llm_chain | parse_taxa\n",
"\n",
"\n",
"def 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",
" )"
]
"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",
@@ -470,40 +190,7 @@
"id": "f1f97ea4-53e5-4f55-8d73-b5b2234a47d9",
"metadata": {},
"outputs": [],
"source": [
"from langgraph.graph import StateGraph\n",
"\n",
"graph = StateGraph(TaxonomyGenerationState)\n",
"graph.add_node(\"summarize\", map_reduce_chain)\n",
"graph.add_node(\"get_minibatches\", get_minibatches)\n",
"graph.add_node(\"generate_taxonomy\", generate_taxonomy)\n",
"graph.add_node(\"update_taxonomy\", update_taxonomy)\n",
"graph.add_node(\"review_taxonomy\", review_taxonomy)\n",
"\n",
"graph.add_edge(\"summarize\", \"get_minibatches\")\n",
"graph.add_edge(\"get_minibatches\", \"generate_taxonomy\")\n",
"graph.add_edge(\"generate_taxonomy\", \"update_taxonomy\")\n",
"\n",
"\n",
"def 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",
"\n",
"graph.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",
")\n",
"graph.set_finish_point(\"review_taxonomy\")\n",
"\n",
"graph.set_entry_point(\"summarize\")\n",
"app = graph.compile()"
]
"source": ["from langgraph.graph import StateGraph, START\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.set_finish_point(\"review_taxonomy\")\n\ngraph.add_edge(START, \"summarize\")\napp = graph.compile()"]
},
{
"cell_type": "code",
@@ -523,11 +210,7 @@
"output_type": "execute_result"
}
],
"source": [
"from IPython.display import Image\n",
"\n",
"Image(app.get_graph().draw_png())"
]
"source": ["from IPython.display import Image\n\nImage(app.get_graph().draw_png())"]
},
{
"cell_type": "markdown",
@@ -549,51 +232,7 @@
"id": "bcc65649-157f-4848-9ef0-8a9932a98d85",
"metadata": {},
"outputs": [],
"source": [
"from datetime import datetime, timedelta\n",
"\n",
"from langsmith import Client\n",
"\n",
"project_name = \"YOUR PROJECT NAME\" # Update to your own project\n",
"client = Client()\n",
"\n",
"past_week = datetime.now() - timedelta(days=7)\n",
"runs = 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.\n",
"def 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",
" }"
]
"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",
@@ -611,15 +250,7 @@
"id": "900906b5-9264-46a8-ba83-46307f8c25d0",
"metadata": {},
"outputs": [],
"source": [
"from langchain.cache import InMemoryCache\n",
"from 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",
"\n",
"set_llm_cache(InMemoryCache())"
]
"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",
@@ -627,39 +258,7 @@
"id": "c2340177-f40c-407a-8e3e-cb06c2ef09ce",
"metadata": {},
"outputs": [],
"source": [
"# We will randomly sample down to 1K docs to speed things up\n",
"docs = [run_to_doc(run) for run in runs if run.inputs]\n",
"docs = random.sample(docs, min(len(docs), 1000))\n",
"use_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",
"\n",
"stream = 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",
"\n",
"for step in stream:\n",
" node, state = next(iter(step.items()))\n",
" print(node, str(state)[:20] + \" ...\")"
]
"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",
@@ -719,31 +318,7 @@
"output_type": "execute_result"
}
],
"source": [
"from IPython.display import Markdown\n",
"\n",
"\n",
"def 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",
"\n",
"Markdown(format_taxonomy_md(step[\"__end__\"][\"clusters\"][-1]))"
]
"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",
@@ -773,32 +348,7 @@
"id": "8aa8a6f5-f53a-41e5-b09d-c6e8476e5471",
"metadata": {},
"outputs": [],
"source": [
"labeling_prompt = hub.pull(\"wfh/tnt-llm-classify\")\n",
"\n",
"labeling_llm = ChatAnthropic(model=\"claude-3-haiku-20240307\", max_tokens_to_sample=2000)\n",
"labeling_llm_chain = (labeling_prompt | labeling_llm | StrOutputParser()).with_config(\n",
" run_name=\"ClassifyDocs\"\n",
")\n",
"\n",
"\n",
"def 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",
"\n",
"labeling_chain = labeling_llm_chain | parse_labels"
]
"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",
@@ -806,23 +356,7 @@
"id": "59c06eea-ecbf-43af-a292-71816ccd92b8",
"metadata": {},
"outputs": [],
"source": [
"final_taxonomy = step[\"__end__\"][\"clusters\"][-1]\n",
"xml_taxonomy = format_taxonomy(final_taxonomy)\n",
"results = 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\n",
"updated_docs = [{**doc, **category} for doc, category in zip(docs, results)]"
]
"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",
@@ -830,10 +364,7 @@
"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: \")"
]
"source": ["if \"OPENAI_API_KEY\" not in os.environ:\n os.environ[\"OPENAI_API_KEY\"] = getpass(\"Enter your OPENAI_API_KEY: \")"]
},
{
"cell_type": "code",
@@ -841,14 +372,7 @@
"id": "c21f787e-2dcb-49c2-9cc1-5284a1732fbc",
"metadata": {},
"outputs": [],
"source": [
"from langchain_openai import OpenAIEmbeddings\n",
"\n",
"# Consider using other embedding models here too!\n",
"encoder = OpenAIEmbeddings(model=\"text-embedding-3-large\")\n",
"vectors = encoder.embed_documents([doc[\"content\"] for doc in docs])\n",
"embedded_docs = [{**doc, \"embedding\": v} for doc, v in zip(updated_docs, vectors)]"
]
"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",
@@ -877,51 +401,7 @@
]
}
],
"source": [
"import numpy as np\n",
"from sklearn.linear_model import LogisticRegression\n",
"from sklearn.metrics import accuracy_score, f1_score\n",
"from sklearn.model_selection import train_test_split\n",
"from sklearn.utils import class_weight\n",
"\n",
"# Create a dictionary mapping category names to their indices in the taxonomy\n",
"category_to_index = {d[\"name\"]: i for i, d in enumerate(final_taxonomy)}\n",
"category_to_index[\"Other\"] = len(category_to_index)\n",
"# Convert category strings to numeric labels\n",
"labels = [\n",
" category_to_index.get(d[\"category\"], category_to_index[\"Other\"])\n",
" for d in embedded_docs\n",
"]\n",
"\n",
"label_vectors = [d[\"embedding\"] for d in embedded_docs]\n",
"\n",
"X_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\n",
"class_weights = class_weight.compute_class_weight(\n",
" class_weight=\"balanced\", classes=np.unique(y_train), y=y_train\n",
")\n",
"class_weight_dict = dict(enumerate(class_weights))\n",
"\n",
"# Weight the classes to partially handle imbalanced data\n",
"model = LogisticRegression(class_weight=class_weight_dict)\n",
"model.fit(X_train, y_train)\n",
"\n",
"train_preds = model.predict(X_train)\n",
"test_preds = model.predict(X_test)\n",
"\n",
"train_acc = accuracy_score(y_train, train_preds)\n",
"test_acc = accuracy_score(y_test, test_preds)\n",
"train_f1 = f1_score(y_train, train_preds, average=\"weighted\")\n",
"test_f1 = f1_score(y_test, test_preds, average=\"weighted\")\n",
"\n",
"print(f\"Train Accuracy: {train_acc:.3f}\")\n",
"print(f\"Test Accuracy: {test_acc:.3f}\")\n",
"print(f\"Train F1 Score: {train_f1:.3f}\")\n",
"print(f\"Test F1 Score: {test_f1:.3f}\")"
]
"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",
@@ -939,15 +419,7 @@
"id": "c27cbb6b-4d0f-476a-bef3-31ed307ce45f",
"metadata": {},
"outputs": [],
"source": [
"from joblib import dump as jl_dump\n",
"\n",
"categories = list(category_to_index)\n",
"\n",
"# Save the model and categories to a file\n",
"with open(\"model.joblib\", \"wb\") as file:\n",
" jl_dump((model, categories), file)"
]
"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",
@@ -965,24 +437,7 @@
"id": "28f0b88a-b308-4208-b482-6c157357dfc6",
"metadata": {},
"outputs": [],
"source": [
"from joblib import load as jl_load\n",
"from langchain_openai import OpenAIEmbeddings\n",
"\n",
"loaded_model, loaded_categories = jl_load(\"model.joblib\")\n",
"encoder = OpenAIEmbeddings(model=\"text-embedding-3-large\")\n",
"\n",
"\n",
"def get_category_name(predictions):\n",
" return [loaded_categories[pred] for pred in predictions]\n",
"\n",
"\n",
"classifier = (\n",
" RunnableLambda(encoder.embed_documents, encoder.aembed_documents)\n",
" | loaded_model.predict\n",
" | get_category_name\n",
")"
]
"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",
@@ -1000,22 +455,7 @@
"id": "6cdb9d8a-2aa1-4f48-8b23-f311fdf36416",
"metadata": {},
"outputs": [],
"source": [
"client = Client()\n",
"\n",
"past_5_min = datetime.now() - timedelta(minutes=5)\n",
"runs = 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",
")\n",
"docs = [run_to_doc(r) for r in runs]"
]
"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",
@@ -1038,10 +478,7 @@
]
}
],
"source": [
"classes = classifier.invoke([doc[\"content\"] for doc in docs])\n",
"print(classes[:2])"
]
"source": ["classes = classifier.invoke([doc[\"content\"] for doc in docs])\nprint(classes[:2])"]
},
{
"cell_type": "markdown",