vision-agent 0.0.10__py3-none-any.whl → 0.0.11__py3-none-any.whl

vision_agent/data/data.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import uuid
 from pathlib import Path
-from typing import Dict, List, Optional, Union, cast
+from typing import Dict, List, Optional, Union, cast, Callable
 
 import faiss
 import numpy as np
@@ -44,18 +44,25 @@ class DataStore:
         self.lmm = lmm
         return self
 
-    def add_column(self, name: str, prompt: str) -> Self:
+    def add_column(
+        self, name: str, prompt: str, func: Optional[Callable[[str], str]] = None
+    ) -> Self:
         r"""Adds a new column to the DataFrame containing the generated metadata from the LMM.
 
         Args:
             name (str): The name of the column to be added.
             prompt (str): The prompt to be used to generate the metadata.
+            func (Optional[Callable[[Any], Any]]): A Python function to be applied on the output of `lmm.generate`. Defaults to None.
         """
         if self.lmm is None:
             raise ValueError("LMM not set yet")
 
         self.df[name] = self.df["image_paths"].progress_apply(  # type: ignore
-            lambda x: self.lmm.generate(prompt, image=x)
+            lambda x: (
+                func(self.lmm.generate(prompt, image=x))
+                if func
+                else self.lmm.generate(prompt, image=x)
+            )
         )
         return self
 

vision_agent/lmm/lmm.py

@@ -31,10 +31,18 @@ class LLaVALMM(LMM):
     def __init__(self, name: str):
         self.name = name
 
-    def generate(self, prompt: str, image: Optional[Union[str, Path]] = None) -> str:
+    def generate(
+        self,
+        prompt: str,
+        image: Optional[Union[str, Path]] = None,
+        temperature: float = 0.1,
+        max_new_tokens: int = 1500,
+    ) -> str:
         data = {"prompt": prompt}
         if image:
             data["image"] = encode_image(image)
+        data["temperature"] = temperature  # type: ignore
+        data["max_new_tokens"] = max_new_tokens  # type: ignore
         res = requests.post(
             _LLAVA_ENDPOINT,
             headers={"Content-Type": "application/json"},

{vision_agent-0.0.10.dist-info → vision_agent-0.0.11.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vision-agent
-Version: 0.0.10
+Version: 0.0.11
 Summary: Toolset for Vision Agent
 Author: Landing AI
 Author-email: dev@landing.ai

{vision_agent-0.0.10.dist-info → vision_agent-0.0.11.dist-info}/RECORD

@@ -1,11 +1,11 @@
 vision_agent/__init__.py,sha256=CpcsVD1aYoCoOD-LNlHAehJPSZBYTtK6CFjiqXUTMAY,172
 vision_agent/data/__init__.py,sha256=YU-5g3LbEQ6a4drz0RLGTagXMVU2Z4Xr3RlfWE-R0jU,46
-vision_agent/data/data.py,sha256=iSLmxPzswQ0g1TOUgBDWBbpF6X8PSAk15oGQxQjhX-o,4833
+vision_agent/data/data.py,sha256=7hM_YprDDgLqreg_qsc-udEUe8cfJQ3-xg6BlVNBG0k,5154
 vision_agent/emb/__init__.py,sha256=YmCkGrJBtXb6X6Z3lnKiFoQYKXMgHMJp8JJyMLVvqcI,75
 vision_agent/emb/emb.py,sha256=la9lhEzk7jqUCjYYQ5oRgVNSnC9_EJBJIpE_B9c6PJo,1375
 vision_agent/lmm/__init__.py,sha256=I8mbeNUajTfWVNqLsuFQVOaNBDlkIhYp9DFU8H4kB7g,51
-vision_agent/lmm/lmm.py,sha256=jPk6rRKUNfKN8P-ShCDosfCtfGkj3aU74x1U8vrZ7yg,2695
-vision_agent-0.0.10.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-vision_agent-0.0.10.dist-info/METADATA,sha256=-6Oyb1M-BjjGhonR0vCvXkEwjIJdDI-phYKfjcQ3lA0,3795
-vision_agent-0.0.10.dist-info/WHEEL,sha256=7Z8_27uaHI_UZAc4Uox4PpBhQ9Y5_modZXWMxtUi4NU,88
-vision_agent-0.0.10.dist-info/RECORD,,
+vision_agent/lmm/lmm.py,sha256=pNUs5XUKUubksFEK_58euE9HPKoPV30mRKtDBEb93MQ,2918
+vision_agent-0.0.11.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+vision_agent-0.0.11.dist-info/METADATA,sha256=CDuOoFqKfzB0F5a_mWRaLi9kvskMouD2gDygdVyxbsk,3795
+vision_agent-0.0.11.dist-info/WHEEL,sha256=7Z8_27uaHI_UZAc4Uox4PpBhQ9Y5_modZXWMxtUi4NU,88
+vision_agent-0.0.11.dist-info/RECORD,,