Skip to content

Commit

Permalink
images wandb logging
Browse files Browse the repository at this point in the history
  • Loading branch information
samshapley committed Sep 2, 2023
1 parent d7b3dc9 commit 199cae0
Show file tree
Hide file tree
Showing 2 changed files with 9 additions and 2 deletions.
3 changes: 2 additions & 1 deletion tools/images.py
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,8 @@ def generate_image(contents_json):
outputs={"caption": prompt},
parent=globals.llm_span,
status=status,
metadata={"error_message": str(e)}
metadata={"error_message": str(e)},
images=[(img, prompt)]
)

return prompt
8 changes: 7 additions & 1 deletion wandb_logging.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
# wandb_logging.py
import wandb
from wandb.sdk.data_types.trace_tree import Trace
import time
from PIL import Image
import globals

def wandb_log_llm(data, model, temperature, parent):
Expand Down Expand Up @@ -45,7 +47,7 @@ def wandb_log_llm(data, model, temperature, parent):



def wandb_log_tool(tool_name, start_time_ms, inputs, outputs, parent, status="success", metadata=None):
def wandb_log_tool(tool_name, start_time_ms, inputs, outputs, parent, status="success", metadata=None, images=None):
time.sleep(1)
end_time_ms = round(time.time() * 1000)

Expand All @@ -69,6 +71,10 @@ def wandb_log_tool(tool_name, start_time_ms, inputs, outputs, parent, status="su
outputs=outputs
)

# Log images if any
if images is not None:
wandb.log({"images": [wandb.Image(img) for img, img_caption in images]})

# update the parent span's end time
parent._span.end_time_ms = end_time_ms

Expand Down

0 comments on commit 199cae0

Please sign in to comment.