new: use a dataclass to store the type of tasks

Signed-off-by: thiswillbeyourgithub <26625900+thiswillbeyourgithub@users.noreply.github.com>

Author: thiswillbeyourgithub

wdoc/__main__.py

@@ -20,7 +20,8 @@
 from wdoc.utils import logger as importedlogger  # make sure to setup the logs first
 from wdoc.wdoc import wdoc
 from wdoc.utils.env import is_out_piped
-from wdoc.utils.misc import get_piped_input, tasks_list
+from wdoc.utils.misc import get_piped_input
+from wdoc.utils.tasks.types import __valid_tasks__
 from wdoc.utils.batch_file_loader import infer_filetype, NoInferrableFiletype
 from typing import Tuple, List, Dict, Any
 import io
@@ -278,13 +279,13 @@ def cli_launcher() -> None:
 
     # turn "wdoc query" into "wdoc --task=query", same for the other tasks
     if "task" not in kwargs:
-        matching_tasks = [t for t in args if t in tasks_list]
+        matching_tasks = [t for t in args if t in __valid_tasks__]
         assert (
             len(matching_tasks) != 0
-        ), f"Found no task in the args: '{args}', wdoc needs one of {tasks_list}"
+        ), f"Found no task in the args: '{args}', wdoc needs one of {__valid_tasks__}"
         assert (
             len(matching_tasks) == 1
-        ), f"Found multiple potential tasks in args: '{args}', wdoc needs one of {tasks_list}"
+        ), f"Found multiple potential tasks in args: '{args}', wdoc needs one of {__valid_tasks__}"
         task = matching_tasks[0]
         logger.debug(f"Moving task '{task}' from args to kwargs")
         args.remove(task)

wdoc/utils/batch_file_loader.py

@@ -26,6 +26,7 @@
 from loguru import logger
 
 from wdoc.utils.env import env, is_out_piped
+from wdoc.utils.tasks.types import wdocTask
 from wdoc.utils.loaders import (
     load_one_doc,
     markdownlink_regex,
@@ -122,7 +123,7 @@ def infer_filetype(path: str) -> str:
 def batch_load_doc(
     llm_name: ModelName,
     filetype: str,
-    task: str,
+    task: wdocTask,
     backend: str,
     n_jobs: int,
     **cli_kwargs,
@@ -297,7 +298,7 @@ def batch_load_doc(
                 "exclude",
             ], "Include or exclude arguments should be reomved at this point"
 
-    if "summar" not in task:
+    if not task.summarize:
         # shuffle the list of files to load to make
         # the hashing progress bar more representative
         to_load = sorted(to_load, key=lambda x: random.random())
@@ -320,7 +321,7 @@ def batch_load_doc(
     for i, h in enumerate(doc_hashes):
         to_load[i]["file_hash"] = doc_hashes[i]
 
-    if "summar" not in task:
+    if not task.summarize:
         # shuffle the list of files again to be random but deterministic:
         # keeping only the digits of each hash, then multiplying by the
         # index of the filetype by size. This makes sure the doc dicts are
@@ -529,7 +530,8 @@ def batch_load_doc(
 
     # smart deduplication before embedding:
     # find the document with the same content_hash, merge their metadata and keep only one
-    if "summar" not in task and len(docs) > 1:
+
+    if task.summarize and len(docs) > 1:
         logger.debug("Deduplicating...")
         logger.debug("Getting all hash")
         content_hash = [d.metadata["content_hash"] for d in docs]

wdoc/utils/loaders/__init__.py

@@ -21,6 +21,7 @@
 from wdoc.utils.env import env
 from wdoc.utils.errors import MissingDocdictArguments, TimeoutPdfLoaderError
 from wdoc.utils.loaders.shared import get_url_title
+from wdoc.utils.tasks.types import wdocTask
 from wdoc.utils.misc import (
     ModelName,
     average_word_length,
@@ -116,7 +117,7 @@ def wrapper(*args, **kwargs) -> Union[List[Document], str]:
 
 @wrapper_load_one_doc
 def load_one_doc(
-    task: str,
+    task: wdocTask,
     llm_name: ModelName,
     temp_dir: Path,
     filetype: str,

wdoc/utils/misc.py

@@ -37,6 +37,7 @@
 
 from wdoc.utils.env import env, is_input_piped, pytest_ongoing
 from wdoc.utils.errors import UnexpectedDocDictArgument
+from wdoc.utils.tasks.types import wdocTask
 
 import lazy_import
 
@@ -140,9 +141,6 @@ def language_detector(text: str) -> None:
 max_token = 10_000_000
 min_lang_prob = 0.50
 
-# list of available tasks
-tasks_list = ["query", "summarize", "parse", "search", "summarize_then_query"]
-
 printed_unexpected_api_keys = [False]  # to print it only once
 
 # loader specific arguments
@@ -713,7 +711,7 @@ def get_tkn_length(
 
 
 def get_splitter(
-    task: str,
+    task: wdocTask,
     modelname: ModelName = DEFAULT_SPLITTER_MODELNAME,
 ) -> "TextSplitter":
     "we don't use the same text splitter depending on the task"
@@ -724,7 +722,7 @@ def get_splitter(
         return text_splitters[task][modelname.original]
 
     # if task is parse but we let the model as testing: assume we want a single super large document with no splitting
-    if task == "parse" and modelname.original == "cliparser/cliparser":
+    if task.parse and modelname.original == "cliparser/cliparser":
         return RecursiveCharacterTextSplitter(
             separators=recur_separator,
             chunk_size=1e7,
@@ -749,7 +747,7 @@ def get_splitter(
         )
 
     # Cap context sizes
-    if task in ["query", "search"] and max_tokens > env.WDOC_MAX_EMBED_CONTEXT:
+    if (task.query or task.search) and max_tokens > env.WDOC_MAX_EMBED_CONTEXT:
         logger.warning(
             f"Capping max_tokens for model {modelname} to WDOC_MAX_EMBED_CONTEXT ({env.WDOC_MAX_EMBED_CONTEXT} instead of {max_tokens}) because in query mode and we can only guess the context size of the embedding model."
         )
@@ -762,27 +760,20 @@ def get_splitter(
 
     model_tkn_length = partial(get_tkn_length, modelname=modelname.original)
 
-    if task in ["query", "search", "parse"]:
+    if task.query or task.search or task.parse:
         text_splitter = RecursiveCharacterTextSplitter(
             separators=recur_separator,
             chunk_size=int(3 / 4 * max_tokens),  # default 4000
             chunk_overlap=500,  # default 200
             length_function=model_tkn_length,
         )
-    elif task in ["summarize_then_query", "summarize"]:
+    elif task.summarize:
         text_splitter = RecursiveCharacterTextSplitter(
             separators=recur_separator,
             chunk_size=int(1 / 2 * max_tokens),
             chunk_overlap=500,
             length_function=model_tkn_length,
         )
-    elif task == "recursive_summary":
-        text_splitter = RecursiveCharacterTextSplitter(
-            separators=recur_separator,
-            chunk_size=int(1 / 4 * max_tokens),
-            chunk_overlap=300,
-            length_function=model_tkn_length,
-        )
     else:
         raise Exception(task)
 

wdoc/utils/retrievers.py

@@ -11,6 +11,7 @@
 from langchain_core.embeddings import Embeddings
 
 from wdoc.utils.env import env
+from wdoc.utils.tasks.types import wdocTask
 from wdoc.utils.misc import cache_dir, get_splitter
 from wdoc.utils.prompts import multiquery_parser, prompts
 from wdoc.utils.customs.compressed_embeddings_cacher import LocalFileStore
@@ -42,7 +43,7 @@ def create_multiquery_retriever(
 
 
 def create_parent_retriever(
-    task: str,
+    task: wdocTask,
     loaded_embeddings: Any,
     loaded_docs: List[Document],
     top_k: int,
@@ -85,7 +86,7 @@ def create_retrievers(
     llm,
     top_k: int,
     relevancy: float,
-    task: str,
+    task: wdocTask,
     loaded_docs: Optional[List[Document]],
 ) -> BaseRetriever:
     """Create and return list of retrievers based on query_retrievers setting."""

wdoc/utils/tasks/parse.py

@@ -12,6 +12,7 @@
 from wdoc.utils.batch_file_loader import batch_load_doc
 from wdoc.utils.logger import debug_exceptions, set_parse_doc_help_md_as_docstring
 from wdoc.utils.misc import DocDict, ModelName
+from wdoc.utils.tasks.types import wdocTask
 
 
 @set_parse_doc_help_md_as_docstring
@@ -77,7 +78,7 @@ def parse_doc(
             cli_kwargs[k] = v
 
     out = batch_load_doc(
-        task="parse",
+        task=wdocTask("parse"),
         filetype=filetype,
         **cli_kwargs,
         **docdict_kwargs,

wdoc/utils/tasks/summarize.py

@@ -13,16 +13,19 @@
 from loguru import logger
 import copy
 from dataclasses import dataclass, asdict
+from functools import partial
+from langchain.text_splitter import RecursiveCharacterTextSplitter
 
 from wdoc.utils.logger import (
     md_printer,
 )
 from wdoc.utils.misc import (
     # debug_chain,
     ModelName,
+    get_model_max_tokens,
     average_word_length,
     check_docs_tkn_length,
-    get_splitter,
+    recur_separator,
     get_tkn_length,
     log_and_time_fn,
     thinking_answer_parser,
@@ -257,6 +260,23 @@ def summarize_documents(
         verbose=llm_verbosity,
     )
 
+    model_tkn_length = partial(get_tkn_length, modelname=model.original)
+    if model.is_testing():
+        max_tokens = 4096
+    else:
+        max_tokens = get_model_max_tokens(model)
+    if max_tokens > env.WDOC_MAX_CHUNK_SIZE:
+        logger.debug(
+            f"Capping max_tokens for model {model.original} to the WDOC_MAX_CHUNK_SIZE value ({env.WDOC_MAX_CHUNK_SIZE} instead of {max_tokens})."
+        )
+        max_tokens = min(max_tokens, env.WDOC_MAX_CHUNK_SIZE)
+    splitter = RecursiveCharacterTextSplitter(
+        separators=recur_separator,
+        chunk_size=int(1 / 4 * max_tokens),
+        chunk_overlap=300,
+        length_function=model_tkn_length,
+    )
+
     # get reading length of the summary
     real_text = "".join([letter for letter in list(summary) if letter.isalpha()])
     sum_reading_length = len(real_text) / average_word_length / wpm
@@ -285,10 +305,6 @@ def summarize_documents(
             assert "- Chunk " not in summary_text, "Found chunk marker"
             assert "- BEFORE RECURSION # " not in summary_text, "Found recursion block"
 
-            splitter = get_splitter(
-                "recursive_summary",
-                modelname=model,
-            )
             summary_docs = [Document(page_content=summary_text)]
             summary_docs = splitter.transform_documents(summary_docs)
             assert summary_docs != relevant_docs

wdoc/utils/tasks/types.py

@@ -0,0 +1,51 @@
+from typing import List
+from dataclasses import dataclass, field
+
+__valid_tasks__: List[str] = [
+    "query",
+    "summarize",
+    "parse",
+    "search",
+    "summarize_then_query",
+]
+
+
+@dataclass
+class wdocTask:
+    original: str
+    query: bool = field(init=False)
+    summarize: bool = field(init=False)
+    parse: bool = field(init=False)
+    search: bool = field(init=False)
+
+    def __post_init__(self):
+        # default values
+        self.query = False
+        self.summarize = False
+        self.parse = False
+        self.search = False
+
+        # checks
+        if isinstance(self.original, wdocTask):
+            self.original = self.original.original
+        self.original = self.original.replace("summary", "summarize")
+        assert (
+            self.original in __valid_tasks__
+        ), f"Received task '{self.original}' is not part of expected tasks: '{__valid_tasks__}'"
+
+        # set the actual properties
+        if "query" in self.original:
+            self.query = True
+        if "summarize" in self.original:
+            self.summarize = True
+        if "search" in self.original:
+            self.search = True
+        if "parse" in self.original:
+            self.parse = True
+
+    def __hash__(self):
+        "necessary for memoizing"
+        return self.original.__hash__()
+
+    def __str__(self) -> str:
+        return self.original