Merge branch 'main' into add_examples

README.md

@@ -24,7 +24,7 @@ import pandas as pd
 
 df = pd.DataFrame({"image_paths": ["image1.png", "image2.png", "image3.png"]})
 ds = va.data.DataStore(df)
-ds = ds.add_lmm(va.lmm.get_model("llava"))
+ds = ds.add_lmm(va.lmm.get_lmm("llava"))
 ds = ds.add_embedder(va.emb.get_embedder("sentence-transformer"))
 
 ds = ds.add_column("descriptions", "Describe this image.")

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "vision-agent"
-version = "0.0.9"
+version = "0.0.10"
 description = "Toolset for Vision Agent"
 authors = ["Landing AI <[email protected]>"]
 readme = "README.md"

vision_agent/lmm/lmm.py

@@ -1,11 +1,16 @@
 import base64
+import logging
 from abc import ABC, abstractmethod
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Union, cast
 
 import requests
 
-from vision_agent.config import BASETEN_API_KEY, BASETEN_URL
+logging.basicConfig(level=logging.INFO)
+
+_LOGGER = logging.getLogger(__name__)
+
+_LLAVA_ENDPOINT = "https://cpvlqoxw6vhpdro27uhkvceady0kvvqk.lambda-url.us-east-2.on.aws"
 
 
 def encode_image(image: Union[str, Path]) -> str:
@@ -39,11 +44,14 @@ def generate(
  data["temperature"] = temperature
  data["max_new_tokens"] = max_new_tokens
  res = requests.post(
- BASETEN_URL,
- headers={"Authorization": f"Api-Key {BASETEN_API_KEY}"},
+ _LLAVA_ENDPOINT,
+ headers={"Content-Type": "application/json"},
  json=data,
  )
- return res.text
+ resp_json: Dict[str, Any] = res.json()
+ if resp_json["statusCode"] != 200:
+ _LOGGER.error(f"Request failed: {resp_json['data']}")
+ return cast(str, resp_json["data"])
 
 
 class OpenAILMM(LMM):