lm-deluge 0.0.3__py3-none-any.whl

lm_deluge/llm_tools/extract.py

@@ -0,0 +1,111 @@
+import io
+import json
+from ..prompt import Conversation
+import asyncio
+from ..client import LLMClient
+from typing import Optional, Any
+from ..util.json import load_json
+
+try:
+    from PIL import Image as PILImage
+except ImportError:
+    PILImage = None
+
+
+async def extract_async(
+    inputs: list[str | Any],
+    schema: Any,
+    client: LLMClient,
+    document_name: Optional[str] = None,
+    object_name: Optional[str] = None,
+    show_progress: bool = True,
+    return_prompts: bool = False,
+):
+    if hasattr(schema, "model_json_schema"):
+        schema_dict = schema.model_json_schema()
+    elif isinstance(schema, dict):
+        schema_dict = schema
+    else:
+        raise ValueError("schema must be a pydantic model or a dict.")
+
+    # warn if json_mode is not True
+    for sp in client.sampling_params:
+        if sp.json_mode is False:
+            print(
+                "Warning: json_mode is False for one or more sampling params. You may get invalid output."
+            )
+            break
+    # check_schema(schema_dict) -- figure out later
+    if document_name is None:
+        document_name = "text"
+    if object_name is None:
+        object_name = ""
+    else:
+        object_name += " "
+
+    text_only_prompt = (
+        f"Given the following {document_name}, extract the {object_name}information "
+        + "from it according to the following JSON schema:\n\n```json\n"
+        + json.dumps(schema_dict, indent=2)
+        + "```\n\nHere is the {document_name}:\n\n```\n{<<__REPLACE_WITH_TEXT__>>}\n```"
+        + "Return the extracted information as JSON, no explanation required. "
+        + f"If the {document_name} seems to be totally irrelevant to the schema (not just incomplete), you may return a JSON object "
+        + 'like `{"error": "The document is not relevant to the schema."}`.'
+    )
+
+    image_only_prompt = (
+        f"Given the attached {document_name} image, extract the {object_name}information "
+        + "from it according to the following JSON schema:\n\n```json\n"
+        + json.dumps(schema_dict, indent=2)
+        + "Return the extracted information as JSON, no explanation required. "
+        + f"If the {document_name} seems to be totally irrelevant to the schema (not just incomplete), you may return a JSON object "
+        + 'like `{"error": "The document is not relevant to the schema."}`.'
+    )
+
+    prompts = []
+    for input in inputs:
+        if isinstance(input, str):
+            prompts.append(
+                text_only_prompt.replace("{<<__REPLACE_WITH_TEXT__>>}", input)
+            )
+        elif PILImage is not None and isinstance(input, PILImage.Image):
+            buffer = io.BytesIO()
+            input.save(buffer, format="PNG")
+            prompts.append(
+                Conversation.user(text=image_only_prompt, image=buffer.getvalue())
+            )
+        else:
+            raise ValueError("inputs must be a list of strings or PIL images.")
+
+    if return_prompts:
+        return prompts
+
+    resps = await client.process_prompts_async(prompts, show_progress=show_progress)
+    completions = [
+        load_json(resp.completion) if (resp and resp.completion) else None
+        for resp in resps
+    ]
+
+    return completions
+
+
+def extract(
+    inputs: list[str | Any],
+    schema: Any,
+    client: LLMClient,
+    document_name: Optional[str] = None,
+    object_name: Optional[str] = None,
+    show_progress: bool = True,
+    return_prompts: bool = False,
+):
+    return asyncio.run(
+        extract_async(
+            inputs,
+            schema,
+            client,
+            document_name,
+            object_name,
+            show_progress,
+            return_prompts,
+        )
+    )

lm_deluge/llm_tools/score.py

@@ -0,0 +1,71 @@
+from ..client import LLMClient, APIResponse
+from ..util.logprobs import extract_prob
+
+# def extract_prob_yes(logprobs: list[dict]):
+#     """
+#     Extract the log probability of the token "yes" from the logprobs object.
+#     Since we can't rely on "yes" and "no" both being in the top_logprobs,
+#     we do the following:
+#     - if token is "yes", return p(yes)
+#     - if token is "no", return 1 - p(no)
+#     - otherwise, return 0.5
+#     """
+#     # use regexp to keep only alpha characters
+#     top_token = logprobs[0]["token"].lower()
+#     top_token = re.sub(r"[^a-z]", "", top_token)
+#     if top_token == "yes":
+#         return np.exp(logprobs[0]["logprob"])
+#     elif top_token == "no":
+#         return 1 - np.exp(logprobs[0]["logprob"])
+#     else:
+#         return 0.5
+
+
+def score_llm(
+    scoring_prompt_template: str,
+    inputs: list[tuple | list | dict],  # to format the template
+    scoring_model: LLMClient,
+    return_probabilities: bool,
+    yes_token: str = "yes",
+) -> list[bool | None] | list[float | None]:
+    if return_probabilities:
+        if not hasattr(scoring_model, "logprobs") or not scoring_model.logprobs:
+            raise ValueError(
+                "return_probabilities=True requires scoring_model to have logprobs=True. "
+                "you may need to upgrade lm_deluge to have access to this option."
+            )
+
+    if scoring_prompt_template is None:
+        raise ValueError("scoring_prompt must be provided.")
+
+    scoring_prompts = []
+    for inp in inputs:
+        if isinstance(inp, dict):
+            scoring_prompt = scoring_prompt_template.format(**inp)
+        elif isinstance(inp, tuple) or isinstance(inp, list):
+            scoring_prompt = scoring_prompt_template.format(*inp)
+        else:
+            raise ValueError("inputs must be a list of tuples, lists, or dicts.")
+        scoring_prompts.append(scoring_prompt)
+
+    resps: list[APIResponse] = scoring_model.process_prompts_sync(  # pyright: ignore
+        prompts=scoring_prompts,
+        show_progress=False,
+    )
+
+    if return_probabilities:
+        logprobs_list = [resp.logprobs for resp in resps]
+        scores = [
+            extract_prob(yes_token, logprobs, use_complement=True)
+            if logprobs is not None
+            else None
+            for logprobs in logprobs_list
+        ]
+    else:
+        completions = [resp.completion for resp in resps]
+        scores = [
+            yes_token.lower().strip() in c.lower() if c is not None else None
+            for c in completions
+        ]
+
+    return scores  # p(yes) or bool yes/no

lm_deluge/llm_tools/translate.py

@@ -0,0 +1,44 @@
+import asyncio
+from ..client import LLMClient
+
+translation_prompt = (
+    "Translate the following text (enclosed in ```) into English. "
+    "Reply with only the translation. Text:\n\n```\n{}\n\n\nYour translation:"
+)
+
+
+def is_english(text: str, low_memory: bool = True):
+    try:
+        from ftlangdetect import detect  # pyright: ignore
+
+        lang = detect(text.replace("\n", " "), low_memory=low_memory)["lang"]
+        return lang == "en"
+    except ImportError:
+        print(
+            "fasttext-langdetect is recommended to use the translate tool, will assume everything is english"
+        )
+    return True
+
+
+async def translate_async(texts: list[str], client: LLMClient, low_memory: bool = True):
+    to_translate_idxs = [
+        i for i, text in enumerate(texts) if not is_english(text, low_memory=low_memory)
+    ]
+    if len(to_translate_idxs) == 0:
+        return texts
+
+    prompts = [translation_prompt.format(texts[i]) for i in to_translate_idxs]
+    resps = await client.process_prompts_async(prompts)
+    translations = [
+        resp.completion.strip() if (resp and resp.completion is not None) else None
+        for resp in resps
+    ]
+    for i, translation in zip(to_translate_idxs, translations):
+        if translation:
+            texts[i] = translation
+
+    return texts
+
+
+def translate(texts: list[str], client: LLMClient, low_memory: bool = True):
+    return asyncio.run(translate_async(texts, client, low_memory))