speedy-utils 1.1.6__py3-none-any.whl → 1.1.7__py3-none-any.whl

llm_utils/__init__.py

@@ -10,7 +10,6 @@ from .chat_format import (
     transform_messages_to_chatml,
 )
 from .lm.async_lm import AsyncLLMTask, AsyncLM
-from .lm.sync_lm import LM, LLMTask
 
 __all__ = [
     "transform_messages",
@@ -21,10 +20,7 @@ __all__ = [
     "display_conversations",
     "build_chatml_input",
     "format_msgs",
-    # "group_messages_by_len",
-    "LM",
-    "AsyncLM",
     "display_chat_messages_as_html",
-    "LLMTask",
+    "AsyncLM",
     "AsyncLLMTask",
 ]

llm_utils/chat_format/transform.py

@@ -16,9 +16,9 @@ def identify_format(item):
 def _transform_sharegpt_to_chatml(
     item, default_system_message="You are a helpful assistant.", print_msg=False
 ):
-    assert isinstance(
-        item, dict
-    ), "The item is not in the correct format. Please check the format of the item."
+    assert isinstance(item, dict), (
+        "The item is not in the correct format. Please check the format of the item."
+    )
 
     messages = []
     system_msg = item.get("system", "")
@@ -116,16 +116,16 @@ def transform_messages_to_chatml(input_data, input_format="auto"):
         input_data = deepcopy(input_data)
         if isinstance(input_data, list):
             input_format = "chatlm"
-            assert (
-                input_data[0].get("role") is not None
-            ), "The input format is not recognized. Please specify the input format."
+            assert input_data[0].get("role") is not None, (
+                "The input format is not recognized. Please specify the input format."
+            )
         elif isinstance(input_data, dict):
             input_data = _transform_sharegpt_to_chatml(input_data)
             input_format = "sharegpt"
         elif isinstance(input_data, str):
-            assert (
-                "<|im_end|>" in input_data
-            ), "The input format is not recognized. Please specify the input format."
+            assert "<|im_end|>" in input_data, (
+                "The input format is not recognized. Please specify the input format."
+            )
             input_format = "chatlm"
             parts = input_data.split("<|im_end|>")
             input_data = []

llm_utils/group_messages.py

@@ -76,7 +76,7 @@ def group_messages_by_len(
     """
     if messages is None:
         raise ValueError("messages parameter cannot be None")
-    from transformers.models.auto.tokenization_auto import AutoTokenizer # type: ignore
+    from transformers.models.auto.tokenization_auto import AutoTokenizer  # type: ignore
 
     tokenizer = AutoTokenizer.from_pretrained(model_name)
 

llm_utils/lm/async_lm/__init__.py

@@ -1,2 +1,7 @@
+from .async_llm_task import AsyncLLMTask
 from .async_lm import AsyncLM
-from .async_llm_task import AsyncLLMTask
+
+__all__ = [
+    "AsyncLM",
+    "AsyncLLMTask",
+]

llm_utils/lm/async_lm/_utils.py

@@ -48,13 +48,17 @@ def _yellow(t):
     return _color(33, t)
 
 
-TParsed = TypeVar("TParsed", bound=BaseModel)
+# TParsed = TypeVar("TParsed", bound=BaseModel)
 
+InputModelType = TypeVar("InputModelType", bound=BaseModel)
+OutputModelType = TypeVar("OutputModelType", bound=BaseModel)
 
-class ParsedOutput(TypedDict, Generic[TParsed]):
+
+class ParsedOutput(TypedDict, Generic[OutputModelType]):
     messages: List
     completion: Any
-    parsed: TParsed
+    parsed: OutputModelType
+    model_kwargs: Dict[str, Any]
 
 
 # --------------------------------------------------------------------------- #
@@ -185,7 +189,6 @@ __all__ = [
     "Messages",
     "LegacyMsgs",
     "RawMsgs",
-    "TParsed",
     "ParsedOutput",
     "get_tokenizer",
     "inspect_word_probs_async",