add prompt to vlm predicate

bdaiinstitute · May 10, 2024 · 51bf569 · 51bf569
1 parent 499f639
commit 51bf569
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/predicators/structs.py b/predicators/structs.py
@@ -463,6 +463,7 @@ class VLMGroundAtom(GroundAtom):
 
     # NOTE: This subclasses GroundAtom to support VLM predicates and classifiers
     predicate: VLMPredicate
+    prompt: Optional[str] = None
 
     def get_query_str(self, without_type: bool = False) -> str:
         """Get a query string for this ground atom.
@@ -476,6 +477,9 @@ def get_query_str(self, without_type: bool = False) -> str:
                 o.name for o in self.objects) + ")"
         else:
             string = str(self)
+
+        if self.prompt is not None:
+            string += f" [Prompt: {self.prompt}]"
         return string
 
     def holds(self, state: State) -> bool: