Upload 14 files

README.md

@@ -1,202 +1,25 @@
----
-metrics:
-- accuracy
-pipeline_tag: question-answering
-tags:
-- code
----
+# Hi-ToM Dataset
 
-# Model Card for Model ID
+This is the dataset for the paper "Hi-ToM: A Benchmark for Evaluating Higher-Order Theory of Mind Reasoning in Large Language Models".
 
-<!-- Provide a quick summary of what the model is/does. -->
+<img src=media/Picture1.png height=430>
 
-This modelcard aims to be a base template for new models. It has been generated using [this raw template](https://github.com/huggingface/huggingface_hub/blob/main/src/huggingface_hub/templates/modelcard_template.md?plain=1).
+### The `Hi-ToM_data` folder
 
-## Model Details
+Contains ToMh data consisting of story-question pairs and the corresponding answers.
+The names of subfolder branches have the following meanings:
 
-### Model Description
+- `Tell` / `No_Tell`: whether or not the stories contain communications among agents.
+- `MC` / `CoT`: the prompting style. `MC` corresponds to Vanilla Prompting (VP) in the paper, while `CoT` stands for Chain-of-Thought Prompting (CoTP).
+- `length_n`: the story length, i.e. the number of chapters in a story. From 1 to 3.
+- `sample_n`: the numbering of different sample stories.
+- `order_n`: the ToM order of the question. From 0 to 4.
 
-<!-- Provide a longer summary of what this model is. -->
+### The `Hi-ToM_prompt` folder
 
+Contains prompt files that can be directly input to API.
+The data in it are almost the same as `Hi-ToM_data`, except that answers are eliminated.
 
+### Generate new data and prompts
 
-- **Developed by:** [More Information Needed]
-- **Funded by [optional]:** [More Information Needed]
-- **Shared by [optional]:** [More Information Needed]
-- **Model type:** [More Information Needed]
-- **Language(s) (NLP):** [More Information Needed]
-- **License:** [More Information Needed]
-- **Finetuned from model [optional]:** [More Information Needed]
-
-### Model Sources [optional]
-
-<!-- Provide the basic links for the model. -->
-
-- **Repository:** [More Information Needed]
-- **Paper [optional]:** [More Information Needed]
-- **Demo [optional]:** [More Information Needed]
-
-## Uses
-
-<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
-
-### Direct Use
-
-<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
-
-[More Information Needed]
-
-### Downstream Use [optional]
-
-<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
-
-[More Information Needed]
-
-### Out-of-Scope Use
-
-<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
-
-[More Information Needed]
-
-## Bias, Risks, and Limitations
-
-<!-- This section is meant to convey both technical and sociotechnical limitations. -->
-
-[More Information Needed]
-
-### Recommendations
-
-<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
-
-Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
-
-## How to Get Started with the Model
-
-Use the code below to get started with the model.
-
-[More Information Needed]
-
-## Training Details
-
-### Training Data
-
-<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
-
-[More Information Needed]
-
-### Training Procedure 
-
-<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
-
-#### Preprocessing [optional]
-
-[More Information Needed]
-
-
-#### Training Hyperparameters
-
-- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
-
-#### Speeds, Sizes, Times [optional]
-
-<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
-
-[More Information Needed]
-
-## Evaluation
-
-<!-- This section describes the evaluation protocols and provides the results. -->
-
-### Testing Data, Factors & Metrics
-
-#### Testing Data
-
-<!-- This should link to a Dataset Card if possible. -->
-
-[More Information Needed]
-
-#### Factors
-
-<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
-
-[More Information Needed]
-
-#### Metrics
-
-<!-- These are the evaluation metrics being used, ideally with a description of why. -->
-
-[More Information Needed]
-
-### Results
-
-[More Information Needed]
-
-#### Summary
-
-
-
-## Model Examination [optional]
-
-<!-- Relevant interpretability work for the model goes here -->
-
-[More Information Needed]
-
-## Environmental Impact
-
-<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
-
-Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
-
-- **Hardware Type:** [More Information Needed]
-- **Hours used:** [More Information Needed]
-- **Cloud Provider:** [More Information Needed]
-- **Compute Region:** [More Information Needed]
-- **Carbon Emitted:** [More Information Needed]
-
-## Technical Specifications [optional]
-
-### Model Architecture and Objective
-
-[More Information Needed]
-
-### Compute Infrastructure
-
-[More Information Needed]
-
-#### Hardware
-
-[More Information Needed]
-
-#### Software
-
-[More Information Needed]
-
-## Citation [optional]
-
-<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
-
-**BibTeX:**
-
-[More Information Needed]
-
-**APA:**
-
-[More Information Needed]
-
-## Glossary [optional]
-
-<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
-
-[More Information Needed]
-
-## More Information [optional]
-
-[More Information Needed]
-
-## Model Card Authors [optional]
-
-[More Information Needed]
-
-## Model Card Contact
-
-[More Information Needed]
+Run the script `generate_tomh.sh`.

actions.py

@@ -0,0 +1,270 @@
+import numpy as np
+
+
+class Action(object):
+
+    def __init__(self, templates):
+        self.templates = templates
+
+    def render_declarative(self, *args):
+        assert 'declarative' in self.templates and \
+            len(self.templates['declarative']) > 0
+        return np.random.choice(self.templates['declarative']) % args
+
+    def render_interrogative(self, *args):
+        assert 'interrogative' in self.templates and \
+            len(self.templates['interrogative']) > 0, str(self.templates)
+        return np.random.choice(self.templates['interrogative']) % args
+
+
+class ExistBeginning(Action):
+
+    def __init__(self):
+        templates = {
+            'interrogative': [
+                'Where was the %s at the beginning?\t%s',
+                'Where was the %s before?\t%s',
+            ]
+        }
+        super().__init__(templates)
+
+
+class Exist(Action):
+
+    def __init__(self):
+        templates = {
+            'interrogative': [
+                'Where is the %s?\t%s',
+                'Where is the %s located?\t%s',
+            ]
+        }
+        super().__init__(templates)
+
+
+class PlaceAction(Action):
+
+    def __init__(self):
+        templates = {
+            'declarative': [
+                '%s placed the %s in the %s.',
+                '%s put the %s in the %s.',
+            ],
+            'interrogative': [
+                'Where did %s place the %s?\t%s',
+                'Where did %s put the %s?\t%s',
+            ]
+        }
+        super().__init__(templates)
+
+
+class SearchAction(Action):
+
+    def __init__(self):
+        templates = {
+            'declarative': [
+                '%s searched for the %s in the %s.',
+                '%s looked for the %s in the %s.',
+            ],
+            'interrogative': [
+                'Where did %s search for the %s?\t%s',
+                'Where did %s look for the %s?\t%s',
+            ],
+        }
+        super().__init__(templates)
+
+
+class TransportAction(Action):
+
+    def __init__(self):
+        templates = {
+            'declarative': [
+                '%s shifted the %s from the %s to the %s.',
+            ],
+        }
+        super().__init__(templates)
+
+
+class EnterAction(Action):
+
+    def __init__(self):
+        templates = {
+            'declarative': [
+                '%s entered the %s.',
+                '%s came into the %s.',
+            ],
+        }
+        super().__init__(templates)
+
+
+class ExitAction(Action):
+
+    def __init__(self):
+        templates = {
+            'declarative': [
+                '%s exited the %s.',
+                '%s left the %s.',
+                '%s went out of the %s.',
+            ],
+        }
+        super().__init__(templates)
+
+
+class BelieveLocationAction(Action):
+
+    def __init__(self):
+        templates = {
+            'declarative': [
+                '%s thinks the %s is in the %s.',
+                '%s believes the %s is in the %s.',
+            ],
+            'interrogative': [
+                'Where does %s think the %s is?\t%s',
+                'Where does %s believe the %s is?\t%s',
+            ],
+        }
+        super().__init__(templates)
+
+
+class BelieveAgentBelieveLocationAction(Action):
+
+    def __init__(self):
+        templates = {
+            'interrogative': [
+                'Where does %s think that %s believes the %s is?\t%s',
+                'Where does %s believe that %s believes the %s is?\t%s',
+                'Where does %s think that %s thinks the %s is?\t%s',
+                'Where does %s believe that %s thinks the %s is?\t%s',
+            ],
+        }
+        super().__init__(templates)
+
+
+class BelieveAgentSearchLocationAction(Action):
+
+    def __init__(self):
+        templates = {
+            'interrogative': [
+                'Where does %s think that %s looks for the %s?\t%s',
+                'Where does %s believe that %s looks for the %s?\t%s',
+                'Where does %s think that %s searches for the %s?\t%s',
+                'Where does %s believe that %s search for the %s?\t%s',
+            ],
+        }
+        super().__init__(templates)
+
+
+class InformLocationAction(Action):
+
+    def __init__(self):
+        templates = {
+            'declarative': [
+                '%s told %s that the %s is in the %s.',
+                '%s informed %s that the %s is in the %s.',
+            ],
+        }
+        super().__init__(templates)
+        
+####################################################
+#######  Deterministic Actions for New Task  #######
+####################################################
+
+class FirstQ(Action):
+
+    def __init__(self):
+        templates = {
+            'interrogative': [
+                'Where will %s look for the %s?\t%s',
+            ]
+        }
+        super().__init__(templates)
+        
+class SecondQ(Action):
+
+    def __init__(self):
+        templates = {
+            'interrogative': [
+                'Where does %s think that %s searches for the %s?\t%s',
+            ]
+        }
+        super().__init__(templates)
+
+class ZeroQ(Action):
+
+    def __init__(self):
+        templates = {
+            'interrogative': [
+                'Where is the %s really?\t%s',
+            ]
+        }
+        super().__init__(templates)
+        
+class MemoryAction(Action):
+
+    def __init__(self):
+        templates = {
+            'interrogative': [
+                'Where was the %s at the beginning?\t%s',
+            ]
+        }
+        super().__init__(templates)
+
+class LocationAction(Action):
+
+    def __init__(self):
+        templates = {
+            'declarative': [
+                '%s and %s are in the %s.',
+            ]
+        }
+        super().__init__(templates)
+
+class ObjectLocAction(Action):
+
+    def __init__(self):
+        templates = {
+            'declarative': [
+                'The %s is in the %s.',
+            ]
+        }
+        super().__init__(templates)
+        
+class ExitedAction(Action):
+
+    def __init__(self):
+        templates = {
+            'declarative': [
+                '%s exited the %s.',
+            ]
+        }
+        super().__init__(templates)
+
+class MoveAction(Action):
+
+    def __init__(self):
+        templates = {
+            'declarative': [
+                '%s moved the %s to the %s.',
+            ]
+        }
+        super().__init__(templates)
+
+class TellAction(Action):
+
+    def __init__(self):
+        templates = {
+            'declarative': [
+                '%s told %s where the %s is.',
+            ]
+        }
+        super().__init__(templates)
+    
+class EnterAction(Action):
+
+    def __init__(self):
+        templates = {
+            'declarative': [
+                '%s entered the %s.',
+            ]
+        }
+        super().__init__(templates)
+    

clause.py

@@ -0,0 +1,26 @@
+import numpy as np
+
+
+class Clause(object):
+
+    def __init__(self, action):
+
+        # if observers is not None:
+        #     assert 0 not in observers, "Observer IDs must be 1-indexed"
+        # self.observers = observers
+        self.action = action
+
+    def render(self):
+        return self.action.render_declarative()  # + \
+        # ('\t' + ' '.join([str(x) for x in self.observers])
+        # if self.observers is not None else '')
+
+
+class Question(Clause):
+
+    def __init__(self, idx_support, action):
+        self.idx_support = idx_support
+        super().__init__(action)
+
+    def render(self):
+        return self.action.render_interrogative()

create_world.py

@@ -0,0 +1,248 @@
+import numpy as np
+
+SIZE_TINY = 5
+SIZE_SMALL = 10
+SIZE_LARGE = 30
+SIZE_XLARGE = 50
+
+locations = [
+    "attic",
+    "back_yard",
+    "basement",
+    "bathroom",
+    "bedroom",
+    "cellar",
+    "closet",
+    "crawlspace",
+    "den",
+    "dining_room",
+    "front_yard",
+    "garage",
+    "garden",
+    "hall",
+    "hallway",
+    "kitchen",
+    "laundry",
+    "living_room",
+    "lounge",
+    "master_bedroom",
+    "office",
+    "pantry",
+    "patio",
+    "playroom",
+    "porch",
+    "staircase",
+    "study",
+    "sunroom",
+    "TV_room",
+    "workshop",
+]
+
+clothing = [
+    "belt",
+    "boots",
+    "cap",
+    "coat",
+    "dress",
+    "gloves",
+    "hat",
+    "jacket",
+    "jeans",
+    "pajamas",
+    "pants",
+    "raincoat",
+    "scarf",
+    "shirt",
+    "shoes",
+    "skirt",
+    "slacks",
+    "slippers",
+    "socks",
+    "stockings",
+    "suit",
+    "sweater",
+    "sweatshirt",
+    "t-shirt",
+    "tie",
+    "trousers",
+    "underclothes",
+    "underpants",
+    "undershirt",
+]
+
+fruit = [
+    "apple",
+    "banana",
+    "cherry",
+    "grapefruit",
+    "grapes",
+    "lemon",
+    "lime",
+    "melon",
+    "orange",
+    "peach",
+    "pear",
+    "persimmon",
+    "pineapple",
+    "plum",
+    "strawberry",
+    "tangerine",
+    "watermelon",
+]
+
+vegetables = [
+    "asparagus",
+    "beans",
+    "broccoli",
+    "cabbage",
+    "carrot",
+    "celery",
+    "corn",
+    "cucumber",
+    "eggplant",
+    "green_pepper",
+    "lettuce",
+    "onion",
+    "peas",
+    "potato",
+    "pumpkin",
+    "radish",
+    "spinach",
+    "sweet_potato",
+    "tomato",
+    "turnip",
+]
+
+objects = fruit + vegetables
+
+containers = [
+    "box",
+    "pantry",
+    "bathtub",
+    "envelope",
+    "drawer",
+    "bottle",
+    "cupboard",
+    "basket",
+    "crate",
+    "suitcase",
+    "bucket",
+    "container",
+    "treasure_chest",
+]
+
+colors = ['green', 'blue', 'red']
+
+containers = ['_'.join([color, container])
+              for container in containers
+              for color in colors]
+
+names = [
+    "Oliver",
+    "Ethan",
+    "Liam",
+    "Benjamin",
+    "Lucas",
+    "Alexander",
+    "Jacob",
+    "Mason",
+    "William",
+    "Gracie",
+    "James",
+    "Logan",
+    "Owen",
+    "Noah",
+    "Carter",
+    "Nathan",
+    "Jack",
+    "Aiden",
+    "Jackson",
+    "Jayden",
+    "Emma",
+    "Olivia",
+    "Emily",
+    "Sophia",
+    "Ava",
+    "Chloe",
+    "Charlotte",
+    "Abigail",
+    "Amelia",
+    "Ella",
+    "Hannah",
+    "Isabella",
+    "Aria",
+    "Lily",
+    "Mia",
+    "Isla",
+    "Avery",
+    "Elizabeth",
+    "Mila",
+    "Evelyn",
+]
+
+assert len(locations) >= SIZE_LARGE
+assert len(objects) >= SIZE_LARGE
+assert len(containers) >= SIZE_LARGE
+assert len(names) >= SIZE_LARGE
+
+
+def write_world(filepath, locs, objs, conts, nams):
+
+    with open(filepath, 'w') as f:
+
+        f.write('# locations\n')
+
+        for loc in locs:
+
+            f.write('\n')
+            f.write('create %s\n' % loc)
+            f.write('set %s is_thing\n' % loc)
+            f.write('set %s is_location\n' % loc)
+
+        f.write('\n')
+        f.write('# objects\n')
+
+        for obj in objs:
+
+            f.write('\n')
+            f.write('create %s\n' % obj)
+            f.write('set %s is_thing\n' % obj)
+            f.write('set %s is_gettable\n' % obj)
+
+        f.write('\n')
+        f.write('# containers\n')
+
+        for cont in conts:
+
+            f.write('\n')
+            f.write('create %s\n' % cont)
+            f.write('set %s is_thing\n' % cont)
+            f.write('set %s is_container\n' % cont)
+
+        f.write('\n')
+        f.write('# actors\n')
+
+        for nam in nams:
+
+            f.write('\n')
+            f.write('create %s\n' % nam)
+            f.write('set %s is_actor\n' % nam)
+            f.write('set %s is_god\n' % nam)
+
+write_world('world_tiny.txt',
+            np.random.choice(locations, SIZE_TINY, replace=False),
+            np.random.choice(objects, SIZE_TINY, replace=False),
+            np.random.choice(containers, SIZE_TINY, replace=False),
+            np.random.choice(names, SIZE_TINY, replace=False))
+
+write_world('world_small.txt',
+            np.random.choice(locations, SIZE_SMALL, replace=False),
+            np.random.choice(objects, SIZE_SMALL, replace=False),
+            np.random.choice(containers, SIZE_SMALL, replace=False),
+            np.random.choice(names, SIZE_SMALL, replace=False))
+
+write_world('world_large.txt',
+            np.random.choice(locations, SIZE_LARGE, replace=False),
+            np.random.choice(objects, SIZE_LARGE, replace=False),
+            np.random.choice(containers, SIZE_LARGE, replace=False),
+            np.random.choice(names, SIZE_LARGE, replace=False))

dynamic_actions.py

@@ -0,0 +1,369 @@
+import numpy as np
+import random
+from itertools import combinations
+from itertools import permutations
+
+
+class Action(object):
+
+    def __init__(self, templates):
+        self.templates = templates
+
+    def render_declarative(self):
+        assert 'declarative' in self.templates and \
+            len(self.templates['declarative']) > 0
+        return np.random.choice(self.templates['declarative'])
+
+    def render_interrogative(self):
+        assert 'interrogative' in self.templates and \
+            len(self.templates['interrogative']) > 0, str(self.templates)
+        return np.random.choice(self.templates['interrogative'])
+
+
+class ExitAction(Action):
+
+    def __init__(self):
+        templates = {
+            'declarative': [
+                '%s exited the %s.',
+                '%s left the %s.',
+                '%s went out of the %s.',
+            ],
+        }
+        super().__init__(templates)
+
+#########################################
+############### Questions ###############
+#########################################
+
+
+class ZeroQ(Action):
+
+    def __init__(self, oracle, obj):
+
+        fill = (obj, oracle.get_object_container(obj))
+        templates = {
+            'interrogative': [
+                'Question: Where is the %s really?\nAnswer: %s' % fill,
+            ]
+        }
+        super().__init__(templates)
+
+
+class FirstQ(Action):
+
+    def __init__(self, oracle, agent, obj):
+        fill = (agent, obj, oracle.get_first_belief(agent, obj))
+        templates = {
+            'interrogative': [
+                'Question: Where does %s really think the %s is?\nAnswer: %s' % fill,
+            ]
+        }
+        super().__init__(templates)
+
+
+class SecondQ(Action):
+
+    def __init__(self, oracle, a1, a2, obj):
+        fill = (a1, a2, obj, oracle.get_second_belief(a1, a2, obj))
+        templates = {
+            'interrogative': [
+                'Question: Where does %s think %s thinks the %s is?\nAnswer: %s' % fill,
+            ]
+        }
+        super().__init__(templates)
+
+
+class ThirdQ(Action):
+
+    def __init__(self, oracle, a1, a2, a3, obj):
+        fill = (a1, a2, a3, obj, oracle.get_third_belief(a1, a2, a3, obj))
+        templates = {
+            'interrogative': [
+                'Question: Where does %s think %s thinks %s thinks the %s is?\nAnswer: %s' % fill,
+            ]
+        }
+        super().__init__(templates)
+
+
+class FourthQ(Action):
+
+    def __init__(self, oracle, a1, a2, a3, a4, obj):
+        fill = (a1, a2, a3, a4, obj,
+                oracle.get_fourth_belief(a1, a2, a3, a4, obj))
+        templates = {
+            'interrogative': [
+                'Question: Where does %s think %s thinks %s thinks %s thinks the %s is?\nAnswer: %s' % fill,
+            ]
+        }
+        super().__init__(templates)
+
+# class MemoryAction(Action):
+
+#     def __init__(self, oracle_start_state, obj):
+#         fill = (obj, oracle_start_state[obj])
+#         templates = {
+#             'interrogative': [
+#                 'Where was the %s at the beginning?\t%s' % fill,
+#             ]
+#         }
+#         super().__init__(templates)
+
+# class LocationAction(Action):
+#     def __init__(self, oracle, args):
+#         """
+#         Creaters string with args and modifies
+#         oracle in accordance with action.
+#         """
+#         if len(args) == 2:
+#             statement = '%s is in the %s.' % args
+#             a1, loc = args
+#             # may be redundant
+#             oracle.set_location(a1, loc)
+#         else : # 2 people
+#             statement = '%s and %s are in the %s.' % args
+#             a1, a2, loc = args
+#             # may be redundant
+#             oracle.set_location(a1, loc)
+#             oracle.set_location(a2, loc)
+
+#         templates = {
+#             'declarative': [
+#                 statement,
+#             ]
+#         }
+
+#         super().__init__(templates)
+
+
+class ObjectLocAction(Action):
+
+    def __init__(self, oracle, obj, observers):
+        container = oracle.get_object_container(obj)
+        templates = {
+            'declarative': [
+                'The %s is in the %s.' % (obj, container),
+            ]
+        }
+
+        # set first beliefs
+        for observer in observers:
+            oracle.set_first_belief(observer, obj, container)
+
+        # set second beliefs
+        if len(observers) >= 2:
+            for observer1, observer2 in combinations(observers, 2):
+                oracle.set_second_belief(observer1, observer2, obj, container)
+                oracle.set_second_belief(observer2, observer1, obj, container)
+
+        # set third beliefs
+        if len(observers) >= 3:
+            for chosen_observers in combinations(observers, 3):
+                for observer1, observer2, observer3 in permutations(chosen_observers):
+                    oracle.set_third_belief(
+                        observer1, observer2, observer3, obj, container)
+
+        # set fourth beliefs
+        if len(observers) >= 4:
+            for chosen_observers in combinations(observers, 4):
+                for observer1, observer2, observer3, observer4 in permutations(chosen_observers):
+                    oracle.set_fourth_belief(
+                        observer1, observer2, observer3, observer4, obj, container)
+        super().__init__(templates)
+
+
+class ExitedAction(Action):
+
+    def __init__(self, oracle, agent):
+        fill = (agent, oracle.get_location(agent))
+
+        templates = {
+            'declarative': [
+                '%s exited the %s.' % fill,
+            ]
+        }
+        oracle.set_location(agent, None)
+        super().__init__(templates)
+
+
+class MoveAction(Action):
+
+    def __init__(self, oracle, args, observers=None, move=True):
+        agent, obj, container = args
+        if not move:
+            location = oracle.get_container_location(container)
+            templates = {
+                'declarative': [
+                    f'{args[0]} made no movements and stayed in the {location} for 1 minute.',
+                ]
+            }
+
+        else:
+            templates = {
+                'declarative': [
+                    '%s moved the %s to the %s.' % args,
+                ]
+            }
+
+            oracle.set_object_container(obj, container)
+
+            if not observers:
+                observers = []
+            observers.append(agent)
+
+            # set first beliefs
+            for observer in observers:
+                oracle.set_first_belief(observer, obj, container)
+
+            # set second beliefs
+            if len(observers) >= 2:
+                for observer1, observer2 in combinations(observers, 2):
+                    oracle.set_second_belief(
+                        observer1, observer2, obj, container)
+                    oracle.set_second_belief(
+                        observer2, observer1, obj, container)
+
+            # set third beliefs
+            if len(observers) >= 3:
+                for chosen_observers in combinations(observers, 3):
+                    for observer1, observer2, observer3 in permutations(chosen_observers):
+                        oracle.set_third_belief(
+                            observer1, observer2, observer3, obj, container)
+
+            # set fourth beliefs
+            if len(observers) >= 4:
+                for chosen_observers in combinations(observers, 4):
+                    for observer1, observer2, observer3, observer4 in permutations(chosen_observers):
+                        oracle.set_fourth_belief(
+                            observer1, observer2, observer3, observer4, obj, container)
+
+        super().__init__(templates)
+
+
+class PublicTellAction(Action):
+
+    def __init__(self, oracle, speaker, obj, container, listeners=None, believers=None):
+        templates = {
+            'declarative': [
+                '%s publicly claimed that %s is in the %s now.' % (
+                    speaker, obj, container),
+            ]
+        }
+        disbelievers = [
+            listener for listener in listeners if listener not in believers]
+
+        # All listeners would think others believe the claim
+        # for believer in believers:
+        #     for disbeliever in disbelievers:
+        #         oracle.set_second_belief(believer, disbeliever, obj, container)
+        #         oracle.set_second_belief(disbeliever, believer, obj, container)
+
+        # A believer would think speaker also believes the obj is in container, speaker would think his words are trusted
+        for believer in believers:
+            oracle.set_first_belief(believer, obj, container)
+            oracle.set_second_belief(believer, speaker, obj, container)
+            oracle.set_second_belief(speaker, believer, obj, container)
+
+        for disbeliever in disbelievers:
+            oracle.set_second_belief(speaker, disbeliever, obj, container)
+
+        # for listener in listeners:
+        #     # the speaker believes that all the listeners believe him
+        #     oracle.set_second_belief(speaker, listener, obj, container)
+        #     # all listeners know the believers based on the exiting order
+        #     for believer in believers:
+        #         oracle.set_second_belief(listener, believer, obj, container)
+
+        super().__init__(templates)
+
+
+class PrivateTellAction(Action):
+
+    def __init__(self, oracle, speaker, listener, obj, container, trust=True):
+        templates = {
+            'declarative': [
+                '%s privately told %s that the %s is in the %s now.' % (
+                    speaker, listener, obj, container),
+            ]
+        }
+
+        # when the listener has less information (exit the room earlier), he'll trust the speaker
+        if trust:
+            oracle.set_first_belief(listener, obj, container)
+            oracle.set_second_belief(listener, speaker, obj, container)
+        oracle.set_second_belief(speaker, listener, obj, container)
+        super().__init__(templates)
+
+
+class EnterAction(Action):
+
+    def __init__(self, oracle, args, observers=None, no_world_adjust=False):
+        templates = {
+            'declarative': [
+                ', '.join(args[:-2]) + ' and ' + args[-2] +
+                ' entered the ' + args[-1] + '.',
+            ]
+        }
+
+        agents = args[:-1]
+        location = args[-1]
+        if location == 'waiting_room':
+            super().__init__(templates)
+            return
+        for agent in agents:
+            oracle.set_location(agent, location)
+        objs = oracle.get_objects_at_location(location)
+        observers = agents
+
+        # agent knows location of everything
+        if not no_world_adjust:
+            for obj in objs:
+                container = oracle.get_object_container(obj)
+                # oracle.set_first_belief(agent, obj, container)
+                # set first beliefs
+                if len(observers) >= 1:
+                    for observer in observers:
+                        oracle.set_first_belief(observer, obj, container)
+
+                # set second beliefs
+                if len(observers) >= 2:
+                    for observer1, observer2 in combinations(observers, 2):
+                        oracle.set_second_belief(
+                            observer1, observer2, obj, container)
+                        oracle.set_second_belief(
+                            observer2, observer1, obj, container)
+
+                # set third beliefs
+                if len(observers) >= 3:
+                    for chosen_observers in combinations(observers, 3):
+                        for observer1, observer2, observer3 in permutations(chosen_observers):
+                            oracle.set_third_belief(
+                                observer1, observer2, observer3, obj, container)
+
+                # set fourth beliefs
+                if len(observers) >= 4:
+                    for chosen_observers in combinations(observers, 4):
+                        for observer1, observer2, observer3, observer4 in permutations(chosen_observers):
+                            oracle.set_fourth_belief(
+                                observer1, observer2, observer3, observer4, obj, container)
+
+        super().__init__(templates)
+
+
+class NoiseAction(Action):
+
+    def __init__(self, agents, containers, objects):
+        animals = ['cat', 'dog', 'monkey', 'mouse']
+        personal_items = ['watch', 'gloves', 'phone']
+        distractors = [
+            f'{random.choice(agents)} saw a {random.choice(animals)}.',
+            f'{random.choice(agents)} lost his {random.choice(personal_items)}.',
+            f'{random.choice(agents)} likes the {random.choice(containers)}.',
+            f'{random.choice(agents)} dislikes the {random.choice(objects)}.',
+        ]
+        templates = {
+            'declarative': [
+                random.choice(distractors)
+            ]
+        }
+        super().__init__(templates)

generate_prompts.py

@@ -0,0 +1,31 @@
+import os
+import sys
+import itertools
+
+
+def main():
+    input_folder = 'data_ToMh'
+    output_folder = 'prompt_ToMh'
+    lengths = [1, 2, 3]
+    orders = [0, 1, 2, 3, 4]
+    prompts = ['CoT', 'MC']
+    tells = ['No_Tell', 'Tell']
+    for tell, prompt, length, order, sample_num in itertools.product(tells, prompts, lengths, orders, range(1, 21)):
+        input_fn = os.path.join(input_folder, tell, prompt, f'length_{length}', f'sample_{sample_num}',
+                                f'order_{order}.txt')
+        output_fn = os.path.join(output_folder, tell, prompt, f'length_{length}', f'sample_{sample_num}',
+                                 f'order_{order}.txt')
+        with open(input_fn, 'r') as file:
+            lines = file.readlines()
+            new_lines = [line for line in lines if line ==
+                         '\n' or line.split()[0] != 'Answer:']
+        if not os.path.exists(os.path.join(output_folder, tell, prompt, f'length_{length}', f'sample_{sample_num}')):
+            os.makedirs(os.path.join(output_folder, tell, prompt,
+                        f'length_{length}', f'sample_{sample_num}'))
+        with open(output_fn, 'w') as file:
+            file.writelines(new_lines)
+
+
+
+if __name__ == "__main__":
+    sys.exit(main())

generate_tasks.py

@@ -0,0 +1,180 @@
+import argparse
+import logging
+import glob
+import numpy as np
+import os
+import sys
+import random
+import itertools
+
+from stringify import stringify
+from tasks import Specify_Tasks
+from utils import is_file, mkdir_p, remove_extension
+from world import World
+
+
+def generate_story_with_specified_chapters(
+    world_paths, output_dir_path, n, noise=0.1, train_noise=False, order=-1, num_chapter=-1, exist_tell_in_story=False, prompt='CoT', exist_answer=False
+):  # prompt is dummy
+    """Generates stories with guarantee that each task is seen n times."""
+    mkdir_p(output_dir_path)
+    n = n[0]
+
+    for world in world_paths:
+
+        w = World()
+        w.load(world)
+        world_name = remove_extension(world)
+
+        # Define task creator and task types
+        task = Specify_Tasks()
+        tasks_per_length = np.array([
+            [('A5', True)],  # 1 chapter
+            [('A5', False), ('A3', True)],  # 2 chapters
+            [('A5', True), ('A3', False), ('A4', True)],  # 3 chapters
+            [('A5', False), ('A3', True),
+             ('A4', False), ('A2', True)],  # 4 chapters
+        ], dtype=object)
+
+        # If order and num_chapter are not specified
+        orders = [0, 1, 2, 3, 4] if order == -1 else [order]
+        num_chapters = [1, 2, 3] if num_chapter == -1 else [num_chapter]
+        modes = ['MC', 'CoT']
+        for length_of_story in num_chapters:
+            # Create folder to contain data
+            folder_name_2 = f'length_{length_of_story}'
+            logging.info("Creating New task in %s..." % folder_name_2)
+
+            for i in range(1, n+1):
+                folder_name_3 = f'sample_{i}'
+                story = task.generate_story_qs_at_end(
+                    w, length_of_story, tasks_per_length[length_of_story -
+                                                         1], num_agents=5,
+                    num_locations=3, statement_noise=noise, order=0, exist_tell_in_story=exist_tell_in_story
+                )  # order = 0 is dummy here.
+                for mode in modes:
+                    folder_name_1 = mode
+                    for order_of_story in orders:
+                        file_name = f'order_{order_of_story}.txt'
+                        os.makedirs(os.path.join(
+                            output_dir_path, folder_name_1, folder_name_2, folder_name_3), exist_ok=True)
+                        path = os.path.join(
+                            output_dir_path, folder_name_1, folder_name_2, folder_name_3, file_name)
+
+                        with open(path, 'w', encoding='utf-8') as f:
+                            if mode == 'MC':
+                                f.write(
+                                    'The following story happens in chronological order. You will be given a multiple-choice question and a note at the end. Directly output the answer without explanation.\n')
+                            else:
+                                f.write(
+                                    'The following story happens in chronological order. You will be given a multiple-choice question and a note at the end. First give step-by-step analysis about the question. Then output the answer.\n')
+                            f.write('Story:\n')
+<<<<<<< HEAD
+                            f.write('\n'.join(stringify(story, exist_answer=exist_answer, order=order_of_story))) #exist_answer is dummy
+                            f.write('\nNote: You should assume the following. (1) A character witnesses everything and every movements before exiting a location. (2) A character A can infer another character B\'s mental state only if A and B have been in the same location, or have private or public interactions. (3) Note that every character tend to lie. A character tend to trust a character that exited the room later than himself. (4) Agents in private communications know that others won\'t hear them, but anyone can hear any public claims. (5) What a character tells others does not affect his actual belief.\n')
+=======
+                            # exist_answer is dummy
+                            f.write(
+                                '\n'.join(stringify(story, exist_answer=exist_answer, order=order_of_story)))
+                            f.write('\nNote: You should assume the following. (1) An agent witnesses everything and every movement before exiting a location. (2) An agent A can infer another agent B\'s mental state only if A and B have been in the same location, or have private or public interactions. (3) Note that every agent tends to lie. What an agent A tells others doesn\'t affect A\'s actual belief. An agent tends to trust an agent that exited the room later than himself. The exit order is known to all agents. (4) Agents in private communications know that others won\'t hear them, but they know that anyone can hear any public claims.\n')
+>>>>>>> 50242d0343261b6c95293fc995711b384ff3c1fe
+
+
+def parse_args(args):
+
+    parser = argparse.ArgumentParser(
+        description='Process command-line arguments.'
+    )
+
+    parser.add_argument(
+        '-w', '--world_path', dest='world_paths', type=is_file, required=True,
+        action='append', help='Path to a world definition file'
+    )
+
+    parser.add_argument(
+        '-o', '--output_dir_path', dest='output_dir_path', type=mkdir_p,
+        default='data_ToMh', help='Output directory path'
+    )
+
+    # parser.add_argument(
+    #     '-b', '--babi_dir_path', dest='babi_dir_path', type=str, required=True,
+    #     help='Path to directory containing the 20 bAbi task train + test data'
+    # )
+
+    parser.add_argument(
+        '-l', '--logging', type=str, default='INFO', metavar='logging',
+        choices=['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'],
+        help='Logging level'
+    )
+
+    parser.add_argument(
+        '-n', '--num_stories', dest='num_stories_choices', type=int,
+        action='append', required=True,
+        help='Number of stories (examples) in a task)'
+    )
+
+    parser.add_argument(
+        '-ptn', '--prob_test_noise', dest='test_noise', type=float,
+        required=True, help='Probability of encountering random noise sentence'
+    )
+
+    parser.add_argument(
+        '-tn', '--train_noise', dest='train_noise', type=bool, default=False,
+        help='Whether or not to include noise at training time'
+    )
+    parser.add_argument(
+        '-ord', '--order', dest='order', type=int, default=-1,
+        help='The range of question orders'
+    )
+    parser.add_argument(
+        '-len', '--length', dest='num_chapter', type=int, default=-1,
+        help='The range of story lengths'
+    )
+    parser.add_argument(
+        '-t', '--tell', dest='exist_tell', type=bool, default=False,
+        help='Whether or not the story has communications between agents'
+    )
+    parser.add_argument(
+        '-p', '--prompt', dest='prompt_type', type=str, default='CoT',
+        choices=['MC', 'CoT'],
+        help='The type of prompt chosen between MC and CoT'
+    )
+    parser.add_argument(
+        '-a', '--answer', dest='exist_answer', type=bool, default=False,
+        help='Whether or not the data has answers'
+    )
+
+    parsed = parser.parse_args(args)
+
+    return parsed
+
+
+def main(args=sys.argv[1:]):
+    """Main function to generate all the story-question pairs."""
+    args = parse_args(args)
+    logging.basicConfig(
+        level=args.logging, format='%(asctime)s\t%(levelname)-8s\t%(message)s'
+    )
+    folder_name = 'Tell/' if args.exist_tell else 'No_Tell/'
+
+    # folder_name += args.prompt_type
+    # output_dir_path = os.path.join(args.output_dir_path, folder_name) if args.exist_answer else os.path.join('prompt_ToMh', folder_name)
+
+    output_dir_path = os.path.join(args.output_dir_path, folder_name)
+
+    generate_story_with_specified_chapters(
+        world_paths=args.world_paths,
+        output_dir_path=output_dir_path,
+        n=args.num_stories_choices,
+        noise=args.test_noise,
+        train_noise=args.train_noise,
+        order=args.order,
+        num_chapter=args.num_chapter,
+        exist_tell_in_story=args.exist_tell,
+        prompt=args.prompt_type,
+        exist_answer=args.exist_answer,
+    )
+
+
+if __name__ == "__main__":
+    sys.exit(main())

oracle.py

@@ -0,0 +1,147 @@
+"""
+The Oracle class keeps track of all object
+and agent locations as well as a map of
+beliefs about object and agent locations.
+"""
+import copy
+
+class LocationMap(object):
+
+    def __init__(self, agents, locations, objects, containers):
+        
+        # Maps agents to their locations.
+        self.locations = {agent : None for agent in agents}
+        
+        # Maps agents to their locations.
+        self.container_locations = {container : None for container in containers}
+        
+        # Maps locations to their containers
+        self.containers = {location : None for location in locations}
+        
+        # Maps containers to the objects they hold
+        self.container_objs = {container : [] for container in containers}
+        
+        # Maps objects to their containers
+        self.obj_containers = {obj : None for obj in objects}
+        
+class MemoryMap(object):
+    
+    def __init__(self, agents, objects):
+        
+        zero_dict = {obj : None for obj in objects}
+        first_dict = {agent : copy.deepcopy(zero_dict) for agent in agents}
+        second_dict = {agent : copy.deepcopy(first_dict) for agent in agents}
+        third_dict = {agent : copy.deepcopy(second_dict) for agent in agents}
+        fourth_dict = {agent : copy.deepcopy(third_dict) for agent in agents}
+        
+        # Dictionary of dictionaries mapping
+        # agents to objects to containers. Represents
+        # agents' belief about location of containers.
+        self.first_belief = copy.deepcopy(first_dict)
+        
+        # Dictionary of dictionaries of dictionaries
+        # mapping agents to direct belief dictionaries.
+        # Represents agents' belief about other agents'
+        # beliefs about location of containers.
+        self.second_belief = copy.deepcopy(second_dict)
+        self.third_belief = copy.deepcopy(third_dict)
+        self.fourth_belief = copy.deepcopy(fourth_dict)
+
+class Oracle(object):
+
+    def __init__(self, agents, locations, objects, containers):
+        self.memory_map = MemoryMap(agents, objects)
+        self.locations = LocationMap(agents, locations, objects, containers)
+        
+    #########################################
+    ################ Beliefs ################
+    #########################################
+    
+    def get_first_belief(self, agent, obj):
+        beliefs = self.memory_map.first_belief
+        return beliefs[agent][obj]
+    
+    def set_first_belief(self, agent, obj, container):
+        beliefs = self.memory_map.first_belief
+        beliefs[agent][obj] = container
+    
+    def get_second_belief(self, a1, a2, obj):
+        second_belief = self.memory_map.second_belief
+        return second_belief[a1][a2][obj]
+            
+    def set_second_belief(self, a1, a2, obj, container):
+        second_belief = self.memory_map.second_belief
+        second_belief[a1][a2][obj] = container
+
+    def get_third_belief(self, a1, a2, a3, obj):
+        third_belief = self.memory_map.third_belief
+        return third_belief[a1][a2][a3][obj]
+            
+    def set_third_belief(self, a1, a2, a3, obj, container):
+        third_belief = self.memory_map.third_belief
+        third_belief[a1][a2][a3][obj] = container
+    
+    def get_fourth_belief(self, a1, a2, a3, a4, obj):
+        fourth_belief = self.memory_map.fourth_belief
+        return fourth_belief[a1][a2][a3][a4][obj]
+            
+    def set_fourth_belief(self, a1, a2, a3, a4, obj, container):
+        fourth_belief = self.memory_map.fourth_belief
+        fourth_belief[a1][a2][a3][a4][obj] = container
+    
+    #########################################
+    ############### Locations ###############
+    #########################################
+    
+    def get_location(self, agent):
+        return self.locations.locations[agent]
+    
+    def set_location(self, agent, location):
+        self.locations.locations[agent] = location
+        
+    def get_containers(self, location):
+        # Returns a list of containers at location
+        return self.locations.containers[location]
+    
+    def set_containers(self, location, containers):
+        # May need to change to move containers bt locs
+        # Containers is a list of containers at location
+        for container in containers:
+            self._set_container_location(container, location)
+        self.locations.containers[location] = containers
+        
+    def get_objects_at_location(self, location):
+        objects = []
+        for container in self.get_containers(location):
+            objects.extend(self.get_container_obj(container))
+        return objects
+       
+    def get_container_location(self, container):
+        return self.locations.container_locations[container]
+    
+    def _set_container_location(self, container, location):
+        self.locations.container_locations[container] = location
+        
+    def get_container_obj(self, container):
+        # get list of objects in container
+        return self.locations.container_objs[container]
+    
+    def _add_container_obj(self, container, obj):
+        self.locations.container_objs[container].append(obj)
+        
+    def _remove_container_obj(self, container, obj):
+        self.locations.container_objs[container].remove(obj)
+    
+    def get_object_container(self, obj):
+        # get container that holds object
+        return self.locations.obj_containers[obj]
+    
+    def set_object_container(self, obj, container):
+        # set container that holds object
+        prev_container = self.get_object_container(obj)
+        if prev_container:
+            self._remove_container_obj(prev_container, obj)
+        self._add_container_obj(container, obj)
+        self.locations.obj_containers[obj] = container
+    
+            

stringify.py

@@ -0,0 +1,47 @@
+import numpy as np
+
+
+def stringify(story, exist_answer=False, order=0): # exist_answer is dummy
+
+    lines = []
+
+    i = 0  # The number of descriptions processed
+    j = 0  # The number of lines output
+    count_order = 0 
+
+    while True:
+
+        
+        if isinstance(story[i], str):
+            line = story[i]
+        else:
+            line = story[i].render()
+            # Capitalize the line
+            line = line[0].upper() + line[1:]
+
+            # Prepend the line number
+            if line.split()[0] != 'Question:' and line.split()[0] != 'Choices:':
+                line = '%d %s' % (i + 1, line)
+            else: # Start with 'Choice'
+                if line.split()[0] == 'Choices:':
+                    lines.append(line)
+                    break
+                else: # Start with 'Question'
+                    if count_order == order:
+                        lines.append(line) 
+                    count_order += 1
+                    i += 1
+                    continue   
+        lines.append(line)
+        # Increment counters
+        i += 1
+
+            # Append supporting lines indices if necessary
+            # if hasattr(story[i], 'idx_support') and story[i].idx_support:
+            #     line += '\t%s' % ' '.join([str(x + 1)
+            #                             for x in story[i].idx_support])
+        
+        if i >= len(story):
+            break
+
+    return lines

tasks.py

@@ -0,0 +1,518 @@
+import numpy as np
+import random
+import copy
+
+from clause import Clause, Question
+from oracle import Oracle
+from dynamic_actions import *
+from collections import defaultdict
+
+
+def sample_question(oracle_start_state, oracle, random_actors, obj, question_idx=0):
+    idx_dummy = [0]
+    a1, a2, a3, a4, _ = random_actors
+    questions = [Question(idx_dummy, ZeroQ(oracle, obj)),
+                 Question(idx_dummy, FirstQ(oracle, a4, obj)),
+                 Question(idx_dummy, SecondQ(oracle, a3, a4, obj)),
+                 Question(idx_dummy, ThirdQ(oracle, a2, a3, a4, obj)),
+                 Question(idx_dummy, FourthQ(oracle, a1, a2, a3, a4, obj))]
+    return questions[question_idx]
+
+#######################################
+############## Chapters ###############
+#######################################
+
+
+def write_A2_chapter(
+        start_state, oracle, obj, location, agent_ids, all_agents, movements=None, exist_tell=False, questions=None
+):
+    a1, a2 = all_agents[agent_ids[0]], all_agents[agent_ids[1]]
+    outsiders = [agent for agent in all_agents if agent not in [a1, a2]]
+    agent_ids = [aid+1 for aid in agent_ids]
+
+    # Pick containers. The first element is the initial container of obj
+    containers = [oracle.get_object_container(obj)]
+    container_candidates = oracle.get_containers(location)[:]
+    container_candidates.remove(containers[0])
+    containers += random.sample(container_candidates, 2)
+
+    # Fill in the chapter
+    chapter = []
+
+    # All selected agents enter the room and see the object
+    chapter.extend([
+        Clause(EnterAction(oracle, (a1, a2, location))),
+        Clause(ObjectLocAction(oracle, obj, [a1, a2])),
+    ])
+
+    # a1
+    chapter.extend([
+        Clause(MoveAction(oracle, (a1, obj, containers[1]), [
+               a2], move=movements[0])),
+        Clause(ExitedAction(oracle, (a1)))
+    ])
+    # a2
+    chapter.extend([
+        Clause(MoveAction(
+            oracle, (a2, obj, containers[2]), None, move=movements[1])),
+        Clause(ExitedAction(oracle, (a2)))
+    ])
+
+    # Everyone enter the waiting room
+    chapter.extend([
+        Clause(EnterAction(oracle, (a1, a2, 'waiting_room')))
+    ])
+
+    # tell actions has 3 different forms
+    if exist_tell:
+        tell_containers = random.sample(oracle.get_containers(location)[:], 2)
+        tell_form = random.choice(
+            range(3)) if outsiders else random.choice(range(2))
+        match tell_form:
+            case 0:
+                chapter.extend([
+                    Clause(PublicTellAction(
+                        oracle, a1, obj, tell_containers[0], listeners=all_agents, believers=outsiders)),
+                    Clause(PrivateTellAction(oracle, a2, a1,
+                           obj, tell_containers[1], trust=True)),
+                ])
+            case 1:
+                chapter.extend([
+                    Clause(PublicTellAction(
+                        oracle, a2, obj, tell_containers[0], listeners=all_agents, believers=[a1] + outsiders)),
+                    Clause(PrivateTellAction(oracle, a1, a2, obj,
+                           tell_containers[1], trust=False)),
+                ])
+            case 2:
+                chapter.extend([
+                    Clause(PrivateTellAction(oracle, a1, random.choice(outsiders),
+                                             obj, tell_containers[0], trust=True))
+                ])
+    return chapter
+
+
+def write_A3_chapter(
+        start_state, oracle, obj, location, agent_ids, all_agents, movements=None, exist_tell=False, questions=None
+):
+    a1, a2, a3 = all_agents[agent_ids[0]
+                            ], all_agents[agent_ids[1]], all_agents[agent_ids[2]]
+    outsiders = [agent for agent in all_agents if agent not in [a1, a2, a3]]
+    agent_ids = [aid+1 for aid in agent_ids]
+
+    # Pick containers. The first element is the initial container of obj
+    containers = [oracle.get_object_container(obj)]
+    container_candidates = oracle.get_containers(location)[:]
+    container_candidates.remove(containers[0])
+    containers += random.sample(container_candidates, 3)
+
+    # Fill in the chapter
+    chapter = []
+
+    # All selected agents enter the room and see the object
+    chapter.extend([
+        Clause(EnterAction(oracle, (a1, a2, a3, location))),
+        Clause(ObjectLocAction(oracle, obj, [a1, a2, a3])),
+    ])
+
+    # a1
+    chapter.extend([
+        Clause(MoveAction(oracle, (a1, obj, containers[1]), [
+               a2, a3], move=movements[0])),
+        Clause(ExitedAction(oracle, (a1)))
+    ])
+    # a2
+    chapter.extend([
+        Clause(MoveAction(oracle, (a2, obj, containers[2]), [
+               a3], move=movements[1])),
+        Clause(ExitedAction(oracle, (a2)))
+    ])
+    # a3
+    chapter.extend([
+        Clause(MoveAction(
+            oracle, (a3, obj, containers[3]), None, move=movements[2])),
+        Clause(ExitedAction(oracle, (a3)))
+    ])
+
+    # Everyone enter the waiting room
+    chapter.extend([
+        Clause(EnterAction(oracle, (a1, a2, a3, 'waiting_room')))
+    ])
+
+    # tell actions has 4 different forms
+    if exist_tell:
+        tell_containers = random.sample(oracle.get_containers(location)[:], 2)
+        tell_form = random.choice(
+            range(4)) if outsiders else random.choice(range(2))
+        match tell_form:
+            case 0:
+                # a2 lies to all, and a3 lies to a2
+                chapter.extend([
+                    Clause(PublicTellAction(
+                        oracle, a2, obj, tell_containers[0], listeners=all_agents, believers=[a1] + outsiders)),
+                    Clause(PrivateTellAction(oracle, a3, a2,
+                           obj, tell_containers[1], trust=True)),
+                ])
+            case 1:
+                # a3 lies to all, and a1 lies to a3
+                chapter.extend([
+                    Clause(PublicTellAction(
+                        oracle, a3, obj, tell_containers[0], listeners=all_agents, believers=[a1, a2] + outsiders)),
+                    Clause(PrivateTellAction(oracle, a1, a3, obj,
+                           tell_containers[1], trust=False)),
+                ])
+            case 2:
+                # a1 lies to all, but a3 tells the true location to an outside agent
+                chapter.extend([
+                    Clause(PublicTellAction(
+                        oracle, a1, obj, tell_containers[0], listeners=all_agents, believers=outsiders)),
+                    Clause(PrivateTellAction(oracle, a3, random.choice(outsiders),
+                           obj, oracle.get_object_container(obj), trust=True))
+                ])
+            case 3:
+                # a2 lies to a3, but a3 tells the true location to an outside agent
+                chapter.extend([
+                    Clause(PrivateTellAction(oracle, a2, a3,
+                           obj, tell_containers[0], trust=False)),
+                    Clause(PrivateTellAction(oracle, a3, random.choice(outsiders),
+                           obj, oracle.get_object_container(obj), trust=True))
+                ])
+    return chapter
+
+
+def write_A4_chapter(
+        start_state, oracle, obj, location, agent_ids, all_agents, movements=None, exist_tell=False, questions=None
+):
+    a1, a2, a3, a4 = all_agents[agent_ids[0]
+                                ], all_agents[agent_ids[1]], all_agents[agent_ids[2]], all_agents[agent_ids[3]]
+    outsiders = [
+        agent for agent in all_agents if agent not in [a1, a2, a3, a4]]
+    agent_ids = [aid+1 for aid in agent_ids]
+
+    # Pick containers. The first element is the initial container of obj
+    containers = [oracle.get_object_container(obj)]
+    container_candidates = oracle.get_containers(location)[:]
+    container_candidates.remove(containers[0])
+    containers += random.sample(container_candidates, 4)
+
+    # Fill in the chapter
+    chapter = []
+
+    # All selected agents enter the room and see the object
+    chapter.extend([
+        Clause(EnterAction(oracle, (a1, a2, a3, a4, location))),
+        Clause(ObjectLocAction(oracle, obj, [a1, a2, a3, a4])),
+    ])
+
+    # a1
+    chapter.extend([
+        Clause(MoveAction(oracle, (a1, obj, containers[1]), [
+               a2, a3, a4], move=movements[0])),
+        Clause(ExitedAction(oracle, (a1)))
+    ])
+    # a2
+    chapter.extend([
+        Clause(MoveAction(oracle, (a2, obj, containers[2]), [
+               a3, a4], move=movements[1])),
+        Clause(ExitedAction(oracle, (a2)))
+    ])
+    # a3
+    chapter.extend([
+        Clause(MoveAction(oracle, (a3, obj, containers[3]), [
+               a4], move=movements[2])),
+        Clause(ExitedAction(oracle, (a3)))
+    ])
+    # a4
+    chapter.extend([
+        Clause(MoveAction(
+            oracle, (a4, obj, containers[4]), None, move=movements[3])),
+        Clause(ExitedAction(oracle, (a4)))
+    ])
+
+    # Everyone enter the waiting room
+    chapter.extend([
+        Clause(EnterAction(oracle, (a1, a2, a3, a4, 'waiting_room')))
+    ])
+
+    # tell actions has 4 different forms
+    if exist_tell:
+        tell_containers = random.sample(oracle.get_containers(location)[:], 2)
+        tell_form = random.choice(
+            range(4)) if outsiders else random.choice(range(2))
+        match tell_form:
+            case 0:
+                # a2 lies to all, and a3 lies to a2
+                chapter.extend([
+                    Clause(PublicTellAction(
+                        oracle, a2, obj, tell_containers[0], listeners=all_agents, believers=[a1] + outsiders)),
+                    Clause(PrivateTellAction(oracle, a4, a3,
+                           obj, tell_containers[1], trust=True)),
+                ])
+            case 1:
+                # a3 lies to all, and a1 lies to a4
+                chapter.extend([
+                    Clause(PublicTellAction(
+                        oracle, a3, obj, tell_containers[0], listeners=all_agents, believers=[a1, a2] + outsiders)),
+                    Clause(PrivateTellAction(oracle, a1, a4, obj,
+                           tell_containers[1], trust=False)),
+                ])
+            case 2:
+                outsider = random.choice(outsiders)
+                # a1 lies to all, but a4 tells the true location to an outside agent
+                chapter.extend([
+                    Clause(PublicTellAction(
+                        oracle, a1, obj, tell_containers[0], listeners=all_agents, believers=outsiders)),
+                    Clause(PrivateTellAction(oracle, a4, outsider,
+                           obj, oracle.get_object_container(obj), trust=True))
+                ])
+            case 3:
+                outsider = random.choice(outsiders)
+                # a2 lies to a3, but a4 tells the true location to an outside agent
+                chapter.extend([
+                    Clause(PrivateTellAction(oracle, a2, a3,
+                           obj, tell_containers[0], trust=False)),
+                    Clause(PrivateTellAction(oracle, a4, outsider,
+                           obj, oracle.get_object_container(obj), trust=True))
+                ])
+    return chapter
+
+
+def write_A5_chapter(
+        start_state, oracle, obj, location, agent_ids, all_agents, movements=None, exist_tell=False, questions=None
+):
+    a1, a2, a3, a4, a5 = all_agents[agent_ids[0]], all_agents[agent_ids[1]
+                                                              ], all_agents[agent_ids[2]], all_agents[agent_ids[3]], all_agents[agent_ids[4]]
+    agent_ids = [aid+1 for aid in agent_ids]
+
+    # Pick containers. The first element is the initial container of obj
+    containers = [oracle.get_object_container(obj)]
+    container_candidates = oracle.get_containers(location)[:]
+    container_candidates.remove(containers[0])
+    containers += random.sample(container_candidates, 4)
+
+    # Fill in the chapter
+    chapter = []
+
+    # All selected agents enter the room and see the object
+    chapter.extend([
+        Clause(EnterAction(oracle, (a1, a2, a3, a4, a5, location))),
+        Clause(ObjectLocAction(oracle, obj, [a1, a2, a3, a4, a5])),
+    ])
+
+    # a1
+    chapter.extend([
+        Clause(MoveAction(oracle, (a1, obj, containers[1]), [
+               a2, a3, a4, a5], move=movements[0])),
+        Clause(ExitedAction(oracle, (a1)))
+    ])
+    # a2
+    chapter.extend([
+        Clause(MoveAction(oracle, (a2, obj, containers[2]), [
+               a3, a4, a5], move=movements[1])),
+        Clause(ExitedAction(oracle, (a2)))
+    ])
+    # a3
+    chapter.extend([
+        Clause(MoveAction(oracle, (a3, obj, containers[3]), [
+               a4, a5], move=movements[2])),
+        Clause(ExitedAction(oracle, (a3)))
+    ])
+    # a4
+    chapter.extend([
+        Clause(MoveAction(oracle, (a4, obj, containers[4]), [
+               a5], move=movements[3])),
+        Clause(ExitedAction(oracle, (a4)))
+    ])
+    # a5
+    chapter.extend([
+        Clause(MoveAction(
+            oracle, (a5, obj, containers[0]), None, move=movements[4])),
+        Clause(ExitedAction(oracle, (a5)))
+    ])
+
+    # Everyone enter the waiting room
+    chapter.extend([
+        Clause(EnterAction(oracle, (a1, a2, a3, a4, a5, 'waiting_room')))
+    ])
+
+    # tell actions has 3 different forms
+    if exist_tell:
+        tell_containers = random.sample(oracle.get_containers(location)[:], 2)
+        tell_form = random.choice(range(3))
+        match tell_form:
+            case 0:
+                # a3 lies to all, and a5 lies to a3
+                chapter.extend([
+                    Clause(PublicTellAction(
+                        oracle, a3, obj, tell_containers[0], listeners=all_agents, believers=[a1, a2])),
+                    Clause(PrivateTellAction(oracle, a5, a3,
+                           obj, tell_containers[1], trust=True)),
+                ])
+            case 1:
+                # a4 lies to all, but a5 tells the true location to a1
+                chapter.extend([
+                    Clause(PublicTellAction(
+                        oracle, a4, obj, tell_containers[0], listeners=all_agents, believers=[a1, a2, a3])),
+                    Clause(PrivateTellAction(oracle, a5, a1, obj,
+                           oracle.get_object_container(obj), trust=True)),
+                ])
+            case 2:
+                # a3 lies a1, and a2 lies to a4
+                chapter.extend([
+                    Clause(PrivateTellAction(oracle, a3, a1,
+                           obj, tell_containers[0], trust=True))
+                ])
+    return chapter
+
+
+#######################################
+############### Tasks #################
+#######################################
+
+class Task(object):
+
+    def __init__(self,
+                 num_questions=5,
+                 exit_prob=1.,
+                 informant_prob=1.,
+                 search_prob=1.,
+                 test_cond='first order'):
+
+        self.num_questions = num_questions
+
+        self.search_prob = search_prob
+
+        self.exit_inform_probs = [1 - exit_prob,
+                                  exit_prob * (1 - informant_prob),
+                                  exit_prob * informant_prob]
+        assert sum(self.exit_inform_probs) == 1
+
+        assert test_cond in ['first order',
+                             'second order',
+                             'reality',
+                             'memory'], \
+            "Invalid test condition: %s" % test_cond
+        self.test_cond = test_cond
+
+    def generate_story(self, world):
+        raise NotImplementedError("Abstract method.")
+
+
+class Specify_Tasks(Task):
+    def generate_story_qs_at_end(
+        self, world, tasks_per_story, tasks, num_agents=5,
+        num_locations=3, statement_noise=0.1, order=0, exist_tell_in_story=False
+    ):
+        """
+        Allows user to specify chapter and question for each task in story.
+
+        :tasks: list with length of tasks per story. Each entry is a string in
+        the set {'tb','fb','sofb'}
+
+        :questions: list with length of tasks per story. Each entry is a string
+        in the set {'memory', 'reality', 'belief', 'search'}
+
+        :statement_noise: probability of encountering noise sentence like 'The
+        dog ran through the kitchen.'
+        """
+
+        # Fetch agents and objects and select a random subset
+        idx_support_dummy = [0]
+        actors = world.get_actors()
+        locations = world.get_locations()
+        objects = world.get_objects()
+        containers = world.get_containers()
+
+        random_actors = np.random.choice(
+            actors, size=num_agents, replace=False
+        )
+        random_locations = np.random.choice(
+            locations, size=num_locations, replace=False
+        )
+        random_objects = np.random.choice(
+            objects, size=num_locations*2, replace=False
+        )
+        random_containers = np.random.choice(
+            containers, size=num_locations*5, replace=False
+        )
+
+        # Create the oracle
+        oracle = Oracle(
+            random_actors, random_locations, random_objects, random_containers
+        )
+
+        # Populate locations in the oracle with containers
+        for i, random_location in enumerate(random_locations):
+            location = random_location
+            containers = random_containers[5*i:5*i+5]
+            oracle.set_containers(location, list(containers))
+            # Two of the containers have objects
+            oracle.set_object_container(
+                random_objects[2*i], containers[0])
+            oracle.set_object_container(
+                random_objects[2*i+1], containers[1])
+
+        # Need start state for memory question
+        start_state = oracle.locations.obj_containers.copy()
+
+        # Create story by task
+        chapters = {'A2': write_A2_chapter,
+                    'A3': write_A3_chapter,
+                    'A4': write_A4_chapter,
+                    'A5': write_A5_chapter}
+        story = []
+        obj_pool = []
+        obj_in_question = None
+
+        for i in range(tasks_per_story):
+            chapter = chapters[tasks[i][0]]
+            location = np.random.choice(random_locations)
+            obj = np.random.choice(oracle.get_objects_at_location(location))
+            # Use the obj in the first chap as the target
+            if i == 0:
+                obj_in_question = obj
+            obj_pool.append(obj)
+            agent_ids = list(range(5))
+            random.shuffle(agent_ids)
+
+            # Randomly choose movements for each agent
+            agent_num = int(tasks[i][0][1])
+            bools = [True, False]
+            movements = [random.choice(bools) for _ in range(agent_num)]
+            exist_tell_in_chapter = tasks[i][1] if exist_tell_in_story else False
+            story.extend(
+                chapter(
+                    start_state, oracle, obj, location, agent_ids, random_actors, movements=movements, exist_tell=exist_tell_in_chapter
+                )
+            )
+
+        # At the end, add noise sentences randomly
+        if statement_noise:
+            noisy_story = []
+            prev_i = 0
+            noise = [i for i
+                     in range(len(story)) if np.random.rand() < statement_noise
+                     ]
+            for i in noise:
+                noisy_story.extend(
+                    story[prev_i:i] +
+                    [Clause(NoiseAction(random_actors,
+                            random_containers, random_objects))]
+                )
+                prev_i = i
+            noisy_story.extend(story[prev_i:])
+
+        # compute questions of all orders
+        questioned_actors = copy.deepcopy(random_actors)
+        random.shuffle(questioned_actors)
+        for idx in range(5):
+            noisy_story.append(
+                sample_question(
+                    start_state, oracle, questioned_actors, obj_in_question, question_idx=idx
+                )
+            )
+
+        # Generate choices of containers
+        choices = ', '.join(f'{chr(65+i)}. {container}' for i,
+                            container in enumerate(random_containers))
+        noisy_story.append('Choices: ' + choices + '\n')
+        return noisy_story

test_azure.py

@@ -0,0 +1,43 @@
+import os
+import openai
+
+def record_progress(filename):
+    with open('progress.txt', 'a') as f:
+        f.write(filename + '\n')
+
+def is_processed(filename):
+    with open('progress.txt', 'r') as f:
+        processed_files = f.read().splitlines()
+    return filename in processed_files
+
+openai.api_type = "azure"
+openai.api_base = "https://openaiserviceforclausaeu.openai.azure.com/"
+openai.api_version = "2023-03-15-preview"
+openai.api_key = os.getenv("OPENAI_API_KEY")
+
+test_dirs = os.listdir("prompt_ToMh")
+for test_dir in test_dirs:
+    test_fns = os.listdir(f"prompt_ToMh/{test_dir}")
+    for test_fn in test_fns:
+        full_path = f"prompt_ToMh/{test_dir}/{test_fn}"
+        if is_processed(full_path):
+            continue
+        print(test_fn)
+        print(f"path: {full_path}")
+        with open(full_path, 'r') as f:
+            input = f.readlines()
+        input = "\n".join([inp.strip() for inp in input])
+        response = openai.ChatCompletion.create(
+          engine="gpt4-32k",
+          messages=[
+            {"role":"system","content":"You are an AI assistant that helps people find information."},
+            {"role":"user","content": input}
+          ],
+          temperature=0,
+          max_tokens=800,
+          top_p=0,
+          frequency_penalty=0,
+          presence_penalty=0,
+          stop=None)
+        print(response)
+        record_progress(full_path)

utils.py

@@ -0,0 +1,44 @@
+from argparse import ArgumentTypeError
+import errno
+import os
+
+
+class Error(Exception):
+    """Base class for exceptions in this module."""
+    pass
+
+
+class InputError(Error):
+    """Exception raised for errors in the input.
+
+    Attributes:
+        expr # input expression in which the error occurred
+        msg  # explanation of the error
+    """
+
+    def __init__(self, expr, msg):
+        self.expr = expr
+        self.msg = msg
+
+
+def is_file(f):
+    try:
+        open(f, 'r')  # return an open file handle
+    except IOError:
+        raise ArgumentTypeError("{0} does not exist".format(f))
+    return f
+
+
+def mkdir_p(path):
+    try:
+        os.makedirs(path)
+    except OSError as exc:  # Python >2.5
+        if exc.errno == errno.EEXIST and os.path.isdir(path):
+            pass
+        else:
+            raise
+    return path
+
+
+def remove_extension(path):
+    return os.path.splitext(os.path.basename(path))[0]

world.py

@@ -0,0 +1,47 @@
+class World(object):
+
+    def __init__(self, world_actions=[], entities={}):
+        self.actions = world_actions
+        self.entities = entities
+
+    def load(self, fname):
+
+        lines = open(fname, 'r').readlines()
+        i = 0
+
+        while i < len(lines):
+            line = lines[i].rstrip('\n')
+            if line != '' and not line.startswith('#'):
+                if line.startswith('create'):
+                    self.entities[line.split(' ')[1]] = {}
+                elif line.startswith('set'):
+                    self.entities[line.split(' ')[1]][line.split(' ')[-1]] = True
+
+            i += 1
+
+    def get_entity(self, predicates):
+
+        if not isinstance(predicates, list):
+            raise InputError(predicates, 'is not a list.')
+
+        return_val = []
+
+        for k in self.entities:
+            if all([predicate in self.entities[k] and
+                    self.entities[k][predicate] is True
+                    for predicate in predicates]):
+                return_val += [k]
+
+        return return_val
+
+    def get_actors(self):
+        return self.get_entity(['is_actor', 'is_god'])
+
+    def get_containers(self):
+        return self.get_entity(['is_thing', 'is_container'])
+
+    def get_locations(self):
+        return self.get_entity(['is_location'])
+
+    def get_objects(self):
+        return self.get_entity(['is_thing', 'is_gettable'])