Fix complaint models

StanfordVL · Mar 5, 2025 · 9befec9 · 9befec9
1 parent 6e8101d
commit 9befec9
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 3 deletions.
diff --git a/bddl/knowledge_base/__init__.py b/bddl/knowledge_base/__init__.py
@@ -11,7 +11,9 @@
   RoomSynsetRequirement, \
   Room, \
   RoomObject, \
-  AttachmentPair
+  AttachmentPair, \
+  ComplaintType, \
+  Complaint
 
 from bddl.knowledge_base.utils import SynsetState
 

diff --git a/bddl/knowledge_base/models.py b/bddl/knowledge_base/models.py
@@ -976,6 +976,7 @@ class Meta:
         pk = "message"
         ordering = ["message"]
 
+    @cached_property
     def objects(self):
         return [complaint.object for complaint in self.complaints]
 

diff --git a/bddl/knowledge_base/processing.py b/bddl/knowledge_base/processing.py
@@ -383,7 +383,7 @@ def create_transitions(self):
                         machine = Synset.get(name=synset_name)
                         transition.machine_synsets.add(machine)
 
-    def create_complaints():
+    def create_complaints(self):
         with open(GENERATED_DATA_DIR / "complaints.json", "r") as f:
             complaints = json.load(f)
 
@@ -401,7 +401,7 @@ def create_complaints():
                 continue
 
             # Create the relevant objects
-            complaint_type, created = Synset.get_or_create(message=complaint_message)
+            complaint_type, created = ComplaintType.get_or_create(message=complaint_message)
             Complaint.create(object=obj, complaint_type=complaint_type, content=complaint_content)
 
     # TODO: Move to cached property on Synset