edsl 0.1.33.dev1__py3-none-any.whl → 0.1.33.dev2__py3-none-any.whl

edsl/questions/templates/multiple_choice/question_presentation.jinja

@@ -0,0 +1,12 @@
+{# Question Presention #}
+{{question_text}}
+{% if use_code %}
+    {%- for option in question_options %}
+{{ loop.index0 }}: {{option}}
+    {% endfor %}
+{% else %}
+    {% for option in question_options %}
+{{option}}
+    {% endfor %}
+{% endif %}
+Only 1 option may be selected.

edsl/questions/templates/numerical/answering_instructions.jinja

@@ -0,0 +1,8 @@
+This question requires a numerical response in the form of an integer or decimal (e.g., -12, 0, 1, 2, 3.45, ...).
+Respond with just your number on a single line.
+If your response is equivalent to zero, report '0'
+If you cannot determine the answer, report 'None'
+
+{% if include_comment %}
+After the answer, put a comment explaining your choice on the next line.
+{% endif %}

edsl/questions/templates/numerical/question_presentation.jinja

@@ -0,0 +1,7 @@
+{{question_text}}
+{% if min_value is not none %}
+    Minimum answer value: {{min_value}}
+{% endif %}
+{% if max_value is not none %}
+    Maximum answer value: {{max_value}}
+{% endif %}

edsl/questions/templates/rank/answering_instructions.jinja

@@ -0,0 +1,11 @@
+{# Answering Instructions #}
+{% if use_code %}
+Please respond only with a comma-separated list of the code of the raked options, with square brackets. E.g., [0, 1, 3]
+{% else %}
+Please respond only with a comma-separated list of the ranked options, with square brackets. E.g., ['Good', 'Bad', 'Ugly']
+{% endif %}
+{% if include_comment %}
+After the answer, you can put a comment explaining your choice on the next line.
+{% endif %}
+
+

edsl/questions/templates/rank/question_presentation.jinja

@@ -0,0 +1,15 @@
+{{question_text}}
+{% if use_code %}
+The options are
+{% for option in question_options %}
+{{ loop.index0 }}: {{option}}
+{% endfor %}
+{% else %}
+The options are:
+{% for option in question_options %}
+{{option}}
+{% endfor %}
+{% endif %}
+{% if num_selections %}
+You can inlcude up to {{num_selections}} options in your answer.
+{% endif %}

edsl/questions/templates/top_k/answering_instructions.jinja

@@ -0,0 +1,8 @@
+{# Answering Instructions #}
+Please respond with valid JSON, formatted like so:
+{% if include_comment %}
+    {"answer": [<put comma-separated list here>], "comment": "<put explanation here>"}
+{% else %}
+    {"answer": [<put comma-separated list here>]}
+{% endif %}
+

edsl/questions/templates/top_k/question_presentation.jinja

@@ -0,0 +1,22 @@
+{{question_text}}
+{% if use_code %}
+    {% for option in question_options %}
+{{ loop.index0 }}: {{option}}
+    {% endfor %}
+{% else %}
+    {% for option in question_options %}
+{{ option }}
+    {% endfor %}    
+{% endif %}
+
+{# Restrictions #}
+{% if min_selections != None and max_selections != None and min_selections == max_selections %}
+You must select exactly {{min_selections}} options.
+{% elif min_selections != None and max_selections != None %}
+Minimum number of options that must be selected: {{min_selections}}.
+Maximum number of options that must be selected: {{max_selections}}.
+{% elif min_selections != None %}
+Minimum number of options that must be selected: {{min_selections}}.
+{% elif max_selections != None %}
+Maximum number of options that must be selected: {{max_selections}}.
+{% endif %}

edsl/questions/templates/yes_no/answering_instructions.jinja

@@ -0,0 +1,6 @@
+{# Answering Instructions #}
+Please reponse with just your answer. 
+
+{% if include_comment %}
+After the answer, you can put a comment explaining your reponse.
+{% endif %}

edsl/questions/templates/yes_no/question_presentation.jinja

@@ -0,0 +1,12 @@
+{# Question Presention #}
+{{question_text}}
+{% if use_code %}
+    {%- for option in question_options %}
+{{ loop.index0 }}: {{option}}
+    {% endfor %}
+{% else %}
+    {% for option in question_options %}
+{{option}}
+    {% endfor %}
+{% endif %}
+Only 1 option may be selected.

edsl/results/Dataset.py

@@ -8,6 +8,7 @@ from typing import Any, Union, Optional
 import numpy as np
 
 from edsl.results.ResultsExportMixin import ResultsExportMixin
+from edsl.results.DatasetTree import Tree
 
 
 class Dataset(UserList, ResultsExportMixin):
@@ -30,6 +31,15 @@ class Dataset(UserList, ResultsExportMixin):
         _, values = list(self.data[0].items())[0]
         return len(values)
 
+    def keys(self):
+        """Return the keys of the first observation in the dataset.
+
+        >>> d = Dataset([{'a.b':[1,2,3,4]}])
+        >>> d.keys()
+        ['a.b']
+        """
+        return [list(o.keys())[0] for o in self]
+
     def __repr__(self) -> str:
         """Return a string representation of the dataset."""
         return f"Dataset({self.data})"
@@ -245,6 +255,16 @@ class Dataset(UserList, ResultsExportMixin):
 
         return Dataset(new_data)
 
+    @property
+    def tree(self):
+        """Return a tree representation of the dataset.
+
+        >>> d = Dataset([{'a':[1,2,3,4]}, {'b':[4,3,2,1]}])
+        >>> d.tree.print_tree()
+        Tree has not been constructed yet.
+        """
+        return Tree(self)
+
     @classmethod
     def example(self):
         """Return an example dataset.

edsl/results/DatasetExportMixin.py

@@ -4,6 +4,7 @@ import base64
 import csv
 import io
 import html
+from typing import Optional
 
 from typing import Literal, Optional, Union, List
 
@@ -41,7 +42,7 @@ class DatasetExportMixin:
         >>> Results.example().relevant_columns(data_type = "flimflam")
         Traceback (most recent call last):
         ...
-        ValueError: No columns found for data type: flimflam. Available data types are: ['agent', 'answer', 'comment', 'model', 'prompt', 'question_options', 'question_text', 'question_type', 'raw_model_response', 'scenario'].
+        ValueError: No columns found for data type: flimflam. Available data types are: ...
         """
         columns = [list(x.keys())[0] for x in self]
         if remove_prefix:
@@ -156,12 +157,13 @@ class DatasetExportMixin:
         iframe_height: int = 200,
         iframe_width: int = 600,
         web=False,
-    ) -> None:
+        return_string: bool = False,
+    ) -> Union[None, str, "Results"]:
         """Print the results in a pretty format.
 
         :param pretty_labels: A dictionary of pretty labels for the columns.
         :param filename: The filename to save the results to.
-        :param format: The format to print the results in. Options are 'rich', 'html', or 'markdown'.
+        :param format: The format to print the results in. Options are 'rich', 'html', 'markdown', or 'latex'.
         :param interactive: Whether to print the results interactively in a Jupyter notebook.
         :param split_at_dot: Whether to split the column names at the last dot w/ a newline.
         :param max_rows: The maximum number of rows to print.
@@ -170,6 +172,9 @@ class DatasetExportMixin:
         :param iframe_height: The height of the iframe.
         :param iframe_width: The width of the iframe.
         :param web: Whether to display the table in a web browser.
+        :param return_string: Whether to return the output as a string instead of printing.
+
+        :return: None if tee is False and return_string is False, the dataset if tee is True, or a string if return_string is True.
 
         Example: Print in rich format at the terminal
 
@@ -253,11 +258,14 @@ class DatasetExportMixin:
 
         >>> r.select('how_feeling').print(format='latex')
         \\begin{tabular}{l}
-        \\toprule
         ...
+        \\end{tabular}
+        <BLANKLINE>
         """
         from IPython.display import HTML, display
         from edsl.utilities.utilities import is_notebook
+        import io
+        import sys
 
         def _determine_format(format):
             if format is None:
@@ -266,7 +274,9 @@ class DatasetExportMixin:
                 else:
                     format = "rich"
             if format not in ["rich", "html", "markdown", "latex"]:
-                raise ValueError("format must be one of 'rich', 'html', or 'markdown'.")
+                raise ValueError(
+                    "format must be one of 'rich', 'html', 'markdown', or 'latex'."
+                )
 
             return format
 
@@ -285,21 +295,24 @@ class DatasetExportMixin:
 
         new_data = list(_create_data())
 
+        # Capture output if return_string is True
+        if return_string:
+            old_stdout = sys.stdout
+            sys.stdout = io.StringIO()
+
+        output = None
+
         if format == "rich":
             from edsl.utilities.interface import print_dataset_with_rich
 
-            print_dataset_with_rich(
+            output = print_dataset_with_rich(
                 new_data, filename=filename, split_at_dot=split_at_dot
             )
-            return self if tee else None
-
-        if format == "markdown":
+        elif format == "markdown":
             from edsl.utilities.interface import print_list_of_dicts_as_markdown_table
 
-            print_list_of_dicts_as_markdown_table(new_data, filename=filename)
-            return self if tee else None
-
-        if format == "latex":
+            output = print_list_of_dicts_as_markdown_table(new_data, filename=filename)
+        elif format == "latex":
             df = self.to_pandas()
             df.columns = [col.replace("_", " ") for col in df.columns]
             latex_string = df.to_latex(index=False)
@@ -309,23 +322,14 @@ class DatasetExportMixin:
                     f.write(latex_string)
             else:
                 print(latex_string)
-
-            return self if tee else None
-
-        if format == "html":
+                output = latex_string
+        elif format == "html":
             from edsl.utilities.interface import print_list_of_dicts_as_html_table
 
             html_source = print_list_of_dicts_as_html_table(
                 new_data, interactive=interactive
             )
 
-            # if download_link:
-            #     from IPython.display import HTML, display
-            #     csv_file = output.getvalue()
-            #     b64 = base64.b64encode(csv_file.encode()).decode()
-            #     download_link = f'<a href="data:file/csv;base64,{b64}" download="my_data.csv">Download CSV file</a>'
-            #     #display(HTML(download_link))
-
             if iframe:
                 iframe = f""""
                 <iframe srcdoc="{ html.escape(html_source) }" style="width: {iframe_width}px; height: {iframe_height}px;"></iframe>
@@ -338,7 +342,18 @@ class DatasetExportMixin:
 
                 view_html(html_source)
 
-            return self if tee else None
+            output = html_source
+
+        # Restore stdout and get captured output if return_string is True
+        if return_string:
+            captured_output = sys.stdout.getvalue()
+            sys.stdout = old_stdout
+            return captured_output or output
+
+        if tee:
+            return self
+
+        return None
 
     def to_csv(
         self,
@@ -501,7 +516,7 @@ class DatasetExportMixin:
 
         return list_of_dicts
 
-    def to_list(self, flatten=False, remove_none=False) -> list[list]:
+    def to_list(self, flatten=False, remove_none=False, unzipped=False) -> list[list]:
         """Convert the results to a list of lists.
 
         :param flatten: Whether to flatten the list of lists.
@@ -596,27 +611,6 @@ class DatasetExportMixin:
         if return_link:
             return filename
 
-    def to_docx(self, filename: Optional[str] = None, separator: str = "\n"):
-        """Export the results to a Word document.
-
-        :param filename: The filename to save the Word document to.
-
-
-        """
-        from docx import Document
-
-        doc = Document()
-        for entry in self:
-            key, values = list(entry.items())[0]
-            doc.add_paragraph(key)
-            line = separator.join(values)
-            doc.add_paragraph(line)
-
-        if filename is not None:
-            doc.save(filename)
-        else:
-            return doc
-
     def tally(
         self, *fields: Optional[str], top_n: Optional[int] = None, output="Dataset"
     ) -> Union[dict, "Dataset"]:

edsl/results/DatasetTree.py

@@ -0,0 +1,145 @@
+from typing import Dict, List, Any, Optional
+from docx import Document
+from docx.shared import Inches, Pt
+from docx.enum.text import WD_ALIGN_PARAGRAPH
+from docx.enum.style import WD_STYLE_TYPE
+
+
+class TreeNode:
+    def __init__(self, key=None, value=None):
+        self.key = key
+        self.value = value
+        self.children = {}
+
+
+class Tree:
+    def __init__(self, data: "Dataset"):
+        d = {}
+        for entry in data:
+            d.update(entry)
+        self.data = d
+        self.root = None
+
+    def unique_values_by_keys(self) -> dict:
+        unique_values = {}
+        for key, values in self.data.items():
+            unique_values[key] = list(set(values))
+        return unique_values
+
+    def construct_tree(self, node_order: Optional[List[str]] = None):
+        # Validate node_order
+        if node_order is None:
+            unique_values = self.unique_values_by_keys()
+            # Sort keys by number of unique values
+            node_order = sorted(
+                unique_values, key=lambda k: len(unique_values[k]), reverse=True
+            )
+        else:
+            if not set(node_order).issubset(set(self.data.keys())):
+                invalid_keys = set(node_order) - set(self.data.keys())
+                raise ValueError(f"Invalid keys in node_order: {invalid_keys}")
+
+        self.root = TreeNode()
+
+        for i in range(len(self.data[node_order[0]])):
+            current = self.root
+            for level in node_order[:-1]:
+                value = self.data[level][i]
+                if value not in current.children:
+                    current.children[value] = TreeNode(key=level, value=value)
+                current = current.children[value]
+
+            leaf_key = node_order[-1]
+            leaf_value = self.data[leaf_key][i]
+            if leaf_value not in current.children:
+                current.children[leaf_value] = TreeNode(key=leaf_key, value=leaf_value)
+
+    def print_tree(
+        self, node: Optional[TreeNode] = None, level: int = 0, print_keys: bool = False
+    ):
+        if node is None:
+            node = self.root
+            if node is None:
+                print("Tree has not been constructed yet.")
+                return
+
+        if node.value is not None:
+            if print_keys and node.key is not None:
+                print("  " * level + f"{node.key}: {node.value}")
+            else:
+                print("  " * level + str(node.value))
+        for child in node.children.values():
+            self.print_tree(child, level + 1, print_keys)
+
+    def to_docx(self, filename: str):
+        doc = Document()
+
+        # Create styles for headings
+        for i in range(1, 10):  # Up to 9 levels of headings
+            style_name = f"Heading {i}"
+            if style_name not in doc.styles:
+                doc.styles.add_style(style_name, WD_STYLE_TYPE.PARAGRAPH)
+
+        # Get or create the 'Body Text' style
+        if "Body Text" not in doc.styles:
+            body_style = doc.styles.add_style("Body Text", WD_STYLE_TYPE.PARAGRAPH)
+        else:
+            body_style = doc.styles["Body Text"]
+
+        body_style.font.size = Pt(11)
+
+        self._add_to_docx(doc, self.root, 0)
+        doc.save(filename)
+
+    def _add_to_docx(self, doc, node: TreeNode, level: int):
+        if node.value is not None:
+            if level == 0:
+                doc.add_heading(str(node.value), level=level + 1)
+            elif node.children:  # If the node has children, it's not the last level
+                para = doc.add_paragraph(str(node.value))
+                para.style = f"Heading {level+1}"
+            else:  # If the node has no children, it's the last level (body text)
+                para = doc.add_paragraph(str(node.value))
+                para.style = "Body Text"
+
+        # Process child nodes (moved outside the if block)
+        for child in node.children.values():
+            self._add_to_docx(doc, child, level + 1)
+
+
+# Example usage (commented out)
+"""
+from edsl.results.Dataset import Dataset
+
+data = Dataset(
+    [
+        {"continent": ["North America", "Asia", "Europe", "North America", "Asia"]},
+        {"country": ["US", "China", "France", "Canada", "Japan"]},
+        {"city": ["New York", "Beijing", "Paris", "Toronto", "Tokyo"]},
+        {"population": [8419000, 21540000, 2161000, 2930000, 13960000]},
+    ]
+)
+
+tree = Tree(data)
+
+try:
+    tree.construct_tree(["continent", "country", "city", "population"])
+    print("Tree without key names:")
+    tree.print_tree()
+    print("\nTree with key names:")
+    tree.print_tree(print_keys=True)
+except ValueError as e:
+    print(f"Error: {e}")
+
+# Demonstrating validation
+try:
+    tree.construct_tree(["continent", "country", "invalid_key"])
+except ValueError as e:
+    print(f"\nValidation Error: {e}")
+
+tree = Tree(data)
+tree.construct_tree(["continent", "country", "city", "population"])
+tree.print_tree(print_keys=True)
+tree.to_docx("tree_structure.docx")
+print("DocX file 'tree_structure.docx' has been created.")
+"""

edsl/results/Result.py

@@ -53,8 +53,8 @@ class Result(Base, UserDict):
 
     >>> import warnings
     >>> warnings.simplefilter("ignore", UserWarning)
-    >>> Result.example().answer
-    {'how_feeling': 'OK', 'how_feeling_comment': 'This is a real survey response from a human.', 'how_feeling_yesterday': 'Great', 'how_feeling_yesterday_comment': 'This is a real survey response from a human.'}
+    >>> Result.example().answer == {'how_feeling_yesterday': 'Great', 'how_feeling': 'OK'}
+    True
 
     Its main data is an Agent, a Scenario, a Model, an Iteration, and an Answer.
     These are stored both in the UserDict and as attributes.
@@ -73,6 +73,8 @@ class Result(Base, UserDict):
         raw_model_response=None,
         survey: Optional["Survey"] = None,
         question_to_attributes: Optional[dict] = None,
+        generated_tokens: Optional[dict] = None,
+        comments_dict: Optional[dict] = None,
     ):
         """Initialize a Result object.
 
@@ -113,6 +115,7 @@ class Result(Base, UserDict):
             "prompt": prompt or {},
             "raw_model_response": raw_model_response or {},
             "question_to_attributes": question_to_attributes,
+            "generated_tokens": generated_tokens or {},
         }
         super().__init__(**data)
         # but also store the data as attributes
@@ -125,6 +128,8 @@ class Result(Base, UserDict):
         self.raw_model_response = raw_model_response or {}
         self.survey = survey
         self.question_to_attributes = question_to_attributes
+        self.generated_tokens = generated_tokens
+        self.comments_dict = comments_dict or {}
 
         self._combined_dict = None
         self._problem_keys = None
@@ -140,7 +145,7 @@ class Result(Base, UserDict):
         else:
             agent_name = self.agent.name
 
-        comments_dict = {k: v for k, v in self.answer.items() if k.endswith("_comment")}
+        # comments_dict = {k: v for k, v in self.answer.items() if k.endswith("_comment")}
         question_text_dict = {}
         question_options_dict = {}
         question_type_dict = {}
@@ -167,11 +172,12 @@ class Result(Base, UserDict):
             "answer": self.answer,
             "prompt": self.prompt,
             "raw_model_response": self.raw_model_response,
-            #            "iteration": {"iteration": self.iteration},
+            "iteration": {"iteration": self.iteration},
             "question_text": question_text_dict,
             "question_options": question_options_dict,
             "question_type": question_type_dict,
-            "comment": comments_dict,
+            "comment": self.comments_dict,
+            "generated_tokens": self.generated_tokens,
         }
 
     def check_expression(self, expression) -> None:
@@ -260,6 +266,26 @@ class Result(Base, UserDict):
             for key, value in subdict.items():
                 yield (index, data_type, key, str(value))
 
+    def leaves(self):
+        leaves = []
+        for question_name, answer in self.answer.items():
+            if not question_name.endswith("_comment"):
+                leaves.append(
+                    {
+                        "question": f"({question_name}): "
+                        + str(
+                            self.question_to_attributes[question_name]["question_text"]
+                        ),
+                        "answer": answer,
+                        "comment": self.answer.get(question_name + "_comment", ""),
+                        "scenario": repr(self.scenario),
+                        "agent": repr(self.agent),
+                        "model": repr(self.model),
+                        "iteration": self.iteration,
+                    }
+                )
+        return leaves
+
     ###############
     # Useful
     ###############
@@ -341,6 +367,7 @@ class Result(Base, UserDict):
                 "raw_model_response", {"raw_model_response": "No raw model response"}
             ),
             question_to_attributes=json_dict.get("question_to_attributes", None),
+            generated_tokens=json_dict.get("generated_tokens", None),
         )
         return result