Spaces:
Sleeping
Sleeping
File size: 8,175 Bytes
399944f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 | from __future__ import annotations
import math
import os
from typing import List, Optional, Sequence, Any, Tuple
from concurrent.futures import ThreadPoolExecutor, as_completed
from kbdebugger.types.ui import ProgressCallback
from rich.progress import track
from kbdebugger.compat.langchain import Document
from kbdebugger.utils import batched
from .sentence_to_qualities import build_sentence_decomposer
from .chunk_to_qualities import build_chunk_decomposer, build_chunk_batch_decomposer
from .types import Qualities, TextDecomposer, BatchTextDecomposer, DecomposeMode
from .logging import save_qualities_json
# ---------------------------------------------------------------------------
# Module-level decomposer singletons
# ---------------------------------------------------------------------------
# These are initialized once at import time to avoid re-loading prompt resources
# and few-shot examples repeatedly inside tight loops.
_sentence_to_qualities_decomposer: TextDecomposer = build_sentence_decomposer()
_chunk_to_qualities_decomposer: TextDecomposer = build_chunk_decomposer()
_chunk_batch_to_qualities_decomposer: BatchTextDecomposer = build_chunk_batch_decomposer()
def decompose(
text: str,
*,
mode: DecomposeMode
) -> Qualities:
"""
Decompose a single input text into "qualities" under the selected mode.
Parameters
----------
text:
Input text: either a single sentence or a larger chunk.
mode:
- DecomposeMode.SENTENCES:
Use when `text` is already sentence-like but may contain
multiple atomic statements that should be split.
Example:
"The cat sat on the mat and looked at the dog."
→ ["The cat sat on the mat.", "The cat looked at the dog."]
- DecomposeMode.CHUNKS:
Use when `text` is a larger paragraph or chunk and you want to
extract key qualities / statements.
Example:
"Cats are great pets. They are independent and curious animals..."
→ ["Cats are great pets.",
"Cats are independent animals.",
"Cats are curious animals."]
Returns
-------
list[str]
A list of short, atomic sentences (qualities).
"""
match mode:
case DecomposeMode.SENTENCES:
return _sentence_to_qualities_decomposer(text)
case DecomposeMode.CHUNKS:
return _chunk_to_qualities_decomposer(text)
case _:
pass
# Defensive: this should never happen with the Enum, but keeps mypy happy
raise ValueError(f"Unsupported DecomposeMode: {mode}")
def _decompose_one_batch(
batch_id: int,
group: List[str],
) -> Tuple[int, List[Qualities]]:
"""
Worker wrapper for parallel batch decomposition.
Returns
-------
(batch_id, batch_results)
batch_id is used to optionally re-order results deterministically.
"""
return batch_id, _chunk_batch_to_qualities_decomposer(group)
def _safe_chunk_batch_to_qualities_decomposer(group: List[str]) -> List[Qualities]:
"""
Safe wrapper around the batched decomposer.
Why this exists
---------------
`ThreadPoolExecutor.map()` will propagate exceptions and stop iteration
on the first failure. For a pipeline stage, it's usually better to be
*best-effort* and preserve output alignment.
Contract
--------
Returns `List[Qualities]` aligned with `group` length:
- one Qualities list per input chunk text
- on failure: returns `[[], [], ...]` (same length as group)
"""
try:
return _chunk_batch_to_qualities_decomposer(group)
except Exception as e: # noqa: BLE001 (intentionally broad in pipeline boundary)
print(f"[decompose_documents] Batch failed (size={len(group)}): {e}")
return [[] for _ in range(len(group))]
# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------
def decompose_documents(
docs: Sequence[Document],
*,
mode: DecomposeMode,
batch_size: int = 5,
use_batch_decomposer: bool = True,
parallel: bool = False,
max_workers: Optional[int] = 2,
progress: Optional[ProgressCallback] = None
) -> Tuple[Qualities, dict]:
"""
Decompose a list of LangChain Documents into a flat list of qualities.
Returns
-------
(qualities, log_payload)
- qualities: flat list of atomic qualities
- log_payload: the same payload that was written to disk
"""
all_qualities: Qualities = []
if not docs:
log_payload = save_qualities_json(
qualities=all_qualities,
mode=mode,
num_input_docs=0,
use_batch_decomposer=use_batch_decomposer,
batch_size=batch_size if use_batch_decomposer else None,
num_batches=0 if use_batch_decomposer else None,
parallel=parallel,
max_workers=max_workers if parallel else None,
)
return all_qualities, log_payload
texts: List[str] = [getattr(doc, "page_content", "") for doc in docs]
# --- Fast path: batched chunk decomposition ---
if mode == DecomposeMode.CHUNKS and use_batch_decomposer:
num_batches = math.ceil(len(texts) / batch_size)
if parallel:
with ThreadPoolExecutor(max_workers=max_workers) as pool:
results_iter = pool.map(
_safe_chunk_batch_to_qualities_decomposer,
batched(texts, batch_size=batch_size),
)
for batch_idx, group_results in track(
enumerate(results_iter),
total=num_batches,
description=(
f"🧷 LLM Decomposer (parallel): paragraphs → qualities "
f"(num_batches={num_batches}, batch size={batch_size})"
),
):
if progress:
progress(
batch_idx + 1, # nicer: 1-based progress
num_batches,
f"🧷 LLM Decomposer (parallel): Processing batch ({batch_idx+1}/{num_batches}) ..."
)
for qualities in group_results:
all_qualities.extend(qualities)
else:
for batch_idx, group in track(
enumerate(batched(texts, batch_size=batch_size)),
total=num_batches,
description=(
f"🧷 LLM Decomposer: paragraphs → qualities "
f"(num_batches={num_batches}, batch size={batch_size})"
),
):
if progress:
progress(
batch_idx + 1,
num_batches,
f"🧷 LLM Decomposer: Processing batch ({batch_idx+1}/{num_batches}) ..."
)
group_results: List[Qualities] = _chunk_batch_to_qualities_decomposer(group)
for qualities in group_results:
all_qualities.extend(qualities)
log_payload = save_qualities_json(
qualities=all_qualities,
mode=mode,
num_input_docs=len(docs),
use_batch_decomposer=True,
batch_size=batch_size,
num_batches=num_batches,
parallel=parallel,
max_workers=max_workers if parallel else None,
)
return all_qualities, log_payload
# --- Default path: one document -> one decompose() call ---
for text in texts:
qualities = decompose(text, mode=mode)
all_qualities.extend(qualities)
log_payload = save_qualities_json(
qualities=all_qualities,
mode=mode,
num_input_docs=len(docs),
use_batch_decomposer=False,
batch_size=None,
num_batches=None,
parallel=False,
max_workers=None,
)
return all_qualities, log_payload
|