Better support for retreiving results and save / load results

Author: yeshaokai

amadeusgpt/configs/Custom_template.yaml

@@ -1,16 +1,11 @@
 keypoint_info:
-  keypoint_file_path:  
-  body_orientation_keypoints:
-    neck: "nose"
-    tail_base: "tail_base"
-    animal_center: "neck"
-  head_orientation_keypoints:
-    nose: "nose"
-    neck: "neck"  
+  keypoint_file_path:
 llm_info:
   keep_last_n_messages: 2  
 object_info:
   load_objects_from_disk: false
 video_info:
   scene_frame_number: 100
   video_file_path: 
+result_info:
+  result_folder: "./results"

amadeusgpt/integration_module_hub.py

@@ -1,4 +1,3 @@
-import glob
 import os
 import pickle
 from typing import Dict, List, Optional

amadeusgpt/main.py

@@ -18,6 +18,8 @@
                                              SelfDebugLLM, VisualLLM)
 from amadeusgpt.integration_module_hub import IntegrationModuleHub
 
+import pickle 
+
 class AMADEUS:
     def __init__(self, config: Dict[str, Any]):
         self.config = config
@@ -81,16 +83,42 @@ def get_analysis(self):
 
     def run_task_program(self, task_program_name: str):
         return self.sandbox.run_task_program(task_program_name)
-
+    
+    def save_results(self, out_folder: str| None = None):
+        """
+        Save the results of the qa message (since it has all the information needed)
+        """
+        if out_folder is None:
+            result_folder = self.sandbox.result_folder
+        else:
+            result_folder = out_folder
+        # make sure it exists
+        os.makedirs(result_folder, exist_ok=True)
+        results = self.sandbox.result_cache
+
+        ret = {}
+        for query in results:
+            ret[query] = results[query].get_serializable()
+
+        # save results to a pickle file
+        with open (os.path.join(result_folder, "results.pickle"), "wb") as f:
+            pickle.dump(ret, f)
+
+    def load_results(self, result_folder: str | None = None ):
+        if result_folder is None:
+            result_folder = self.sandbox.result_folder
+        else:
+            result_folder = result_folder
+        
+        with open (os.path.join(result_folder, "results.pickle"), "rb") as f:
+            results = pickle.load(f)
+        self.sandbox.result_cache = results
+
+    def get_results(self):
+        return self.sandbox.result_cache
 
 if __name__ == "__main__":
     from amadeusgpt.analysis_objects.llm import VisualLLM
     from amadeusgpt.config import Config
-
-
     config = Config("amadeusgpt/configs/EPM_template.yaml")
-
-    amadeus = AMADEUS(config)
-    sandbox = amadeus.sandbox
-    visualLLm = VisualLLM(config)
-    visualLLm.speak(sandbox)
+    amadeus = AMADEUS(config)   

amadeusgpt/programs/mabe_social_mabe.py

@@ -6,10 +6,8 @@
 import traceback
 import typing
 from functools import wraps
-
 import matplotlib.pyplot as plt
 import numpy as np
-
 from amadeusgpt.analysis_objects.analysis_factory import create_analysis
 from amadeusgpt.analysis_objects.event import BaseEvent
 from amadeusgpt.analysis_objects.relationship import Orientation
@@ -18,10 +16,42 @@
                                               INTEGRATION_API_REGISTRY)
 from amadeusgpt.programs.task_program_registry import (TaskProgram,
                                                        TaskProgramLibrary)
+from pathlib import Path
+
+
+class QA_Message(dict):
+    def __init__(self, *args, **kwargs):
+        super(QA_Message, self).__init__(*args, **kwargs)
+
+    def get_masks(self):
+        function_rets = self["function_rets"]
+        # if function_ret is a list of events
+        if (
+            isinstance(function_rets, list)
+            and len(function_rets) > 0
+            and isinstance(function_rets[0], BaseEvent)
+        ):
+            events = function_rets
+            masks = []
+            for event in events:
+                masks.append(event.generate_mask())
+            return np.array(masks)
+        else:
+            raise ValueError("No events found in the function_rets")
+        
+    def get_serializable(self):
+        """
+        Only part of qa messages are serializable.
+        """
+        selected_keys = ['query', 'code', 'chain_of_thought', 'function_rets', 'meta_info']
+        ret = {}
+        for key in selected_keys:
+            ret[key] = self[key]
+        return ret
 
 
 def create_message(query, sandbox):
-    return {
+    return QA_Message({
         "query": query,
         "code": None,
         "chain_of_thought": None,
@@ -32,7 +62,7 @@ def create_message(query, sandbox):
         "out_videos": None,
         "pose_video": None,
         "meta_info": None,
-    }
+    })
 
 
 class SandboxBase:
@@ -192,16 +222,21 @@ def __init__(self, config):
         self.task_program_library = TaskProgramLibrary().get_task_programs()
         self.config = config
         self.messages = []
+        # initialize the code execution namespace with builtins
         self.exec_namespace = {"__builtins__": __builtins__}
         # update_namespace initializes behavior analysis
         self.update_namespace()
         # then we can configure behavior analysis using vlm
         self.meta_info = None
-        self.visual_cache = {}
+        # where llms are stored
         self.llms = {}
         # just easier to pass this around
         self.query = None
         self.matched_modules = []
+        # result cache keeps the qa_message using the query as the key:
+        self.result_cache = {}
+        # configure how to save the results to a result folder
+        self.result_folder = Path(self.config["result_info"].get("result_folder", "./results"))
 
     def configure_using_vlm(self):
         # example meta_info:
@@ -322,9 +357,6 @@ def update_namespace(self):
         # to allow the program to access existing task programs
         self.exec_namespace["task_programs"] = TaskProgramLibrary.get_task_programs()
 
-    def parse_function_results(self, function_rets):
-        pass
-
     def code_execution(self, qa_message):
         # add main function into the namespace
         self.update_namespace()
@@ -388,7 +420,8 @@ def register_llm(self, name, llm):
     def events_to_videos(self, events, function_name):
         behavior_analysis = self.exec_namespace["behavior_analysis"]
         visual_manager = behavior_analysis.visual_manager
-        out_folder = "event_clips"
+        # save video clips to the result folder
+        out_folder = str(self.result_folder)
         os.makedirs(out_folder, exist_ok=True)
         behavior_name = "_".join(function_name.split(" "))
         video_file = self.config["video_info"]["video_file_path"]
@@ -449,9 +482,7 @@ def render_qa_message(self, qa_message):
             qa_message["out_videos"] = self.events_to_videos(
                 function_rets, self.get_function_name_from_string(qa_message["code"])
             )
-
-        else:
-            pass
+     
         qa_message["plots"].extend(plots)
         return qa_message
 
@@ -463,22 +494,25 @@ def llm_step(self, user_query):
             qa_message["meta_info"] = self.meta_info
 
         self.messages.append(qa_message)
+        # there might be better way to set this
         self.query = user_query
         self.llms["code_generator"].speak(self)
-
+        self.result_cache[user_query] = qa_message
         return qa_message
 
     def run_task_program(self, task_program_name):
         """
         Sandbox is also responsible for running task program
         """
         task_program = self.task_program_library[task_program_name]
-        self.query = "run the task program"
+        # there might be better way to set this
+        self.query = task_program_name
         qa_message = create_message(self.query, self)
         qa_message["code"] = task_program["source_code"]
         self.messages.append(qa_message)
         self.code_execution(qa_message)
         qa_message = self.render_qa_message(qa_message)
+        self.result_cache[task_program_name] = qa_message
         return qa_message
 
     def step(self, user_query, number_of_debugs=1):
@@ -503,7 +537,7 @@ def step(self, user_query, number_of_debugs=1):
                 qa_message = self.code_execution(qa_message)
 
         qa_message = self.render_qa_message(qa_message)
-
+        self.result_cache[user_query] = qa_message
         return qa_message