edsl 0.1.32__py3-none-any.whl → 0.1.33__py3-none-any.whl

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", {}),
         )
         return result
 

edsl/results/Results.py

@@ -17,6 +17,7 @@ from edsl.exceptions.results import (
     ResultsInvalidNameError,
     ResultsMutateError,
     ResultsFilterError,
+    ResultsDeserializationError,
 )
 
 from edsl.results.ResultsExportMixin import ResultsExportMixin
@@ -77,6 +78,7 @@ class Results(UserList, Mixins, Base):
         "question_options",
         "question_type",
         "comment",
+        "generated_tokens",
     ]
 
     def __init__(
@@ -108,6 +110,81 @@ class Results(UserList, Mixins, Base):
         if hasattr(self, "_add_output_functions"):
             self._add_output_functions()
 
+    def leaves(self):
+        leaves = []
+        for result in self:
+            leaves.extend(result.leaves())
+        return leaves
+
+    def tree(
+        self,
+        fold_attributes: Optional[List[str]] = None,
+        drop: Optional[List[str]] = None,
+        open_file=True,
+    ) -> dict:
+        """Return the results as a tree."""
+        from edsl.results.tree_explore import FoldableHTMLTableGenerator
+
+        if drop is None:
+            drop = []
+
+        valid_attributes = [
+            "model",
+            "scenario",
+            "agent",
+            "answer",
+            "question",
+            "iteration",
+        ]
+        if fold_attributes is None:
+            fold_attributes = []
+
+        for attribute in fold_attributes:
+            if attribute not in valid_attributes:
+                raise ValueError(
+                    f"Invalid fold attribute: {attribute}; must be in {valid_attributes}"
+                )
+        data = self.leaves()
+        generator = FoldableHTMLTableGenerator(data)
+        tree = generator.tree(fold_attributes=fold_attributes, drop=drop)
+        html_content = generator.generate_html(tree, fold_attributes)
+        import tempfile
+        from edsl.utilities.utilities import is_notebook
+
+        from IPython.display import display, HTML
+
+        if is_notebook():
+            import html
+            from IPython.display import display, HTML
+
+            height = 1000
+            width = 1000
+            escaped_output = html.escape(html_content)
+            # escaped_output = rendered_html
+            iframe = f""""
+            <iframe srcdoc="{ escaped_output }" style="width: {width}px; height: {height}px;"></iframe>
+            """
+            display(HTML(iframe))
+            return None
+
+        with tempfile.NamedTemporaryFile(suffix=".html", delete=False) as f:
+            f.write(html_content.encode())
+            print(f"HTML file has been generated: {f.name}")
+
+            if open_file:
+                import webbrowser
+                import time
+
+                time.sleep(1)  # Wait for 1 second
+                # webbrowser.open(f.name)
+                import os
+
+                filename = f.name
+                webbrowser.open(f"file://{os.path.abspath(filename)}")
+
+            else:
+                return html_content
+
     def code(self):
         raise NotImplementedError
 
@@ -168,7 +245,9 @@ class Results(UserList, Mixins, Base):
         )
 
     def __repr__(self) -> str:
-        return f"Results(data = {self.data}, survey = {repr(self.survey)}, created_columns = {self.created_columns})"
+        import reprlib
+
+        return f"Results(data = {reprlib.repr(self.data)}, survey = {repr(self.survey)}, created_columns = {self.created_columns})"
 
     def _repr_html_(self) -> str:
         from IPython.display import HTML
@@ -290,8 +369,7 @@ class Results(UserList, Mixins, Base):
                 ),
             )
         except Exception as e:
-            print(e)
-            # breakpoint()
+            raise ResultsDeserializationError(f"Error in Results.from_dict: {e}")
         return results
 
     ######################
@@ -395,7 +473,7 @@ class Results(UserList, Mixins, Base):
 
         >>> r = Results.example()
         >>> r.models[0]
-        Model(model_name = 'gpt-4-1106-preview', temperature = 0.5, max_tokens = 1000, top_p = 1, frequency_penalty = 0, presence_penalty = 0, logprobs = False, top_logprobs = 3)
+        Model(model_name = ...)
         """
         return [r.model for r in self.data]
 
@@ -477,39 +555,6 @@ class Results(UserList, Mixins, Base):
         )
         return sorted(list(all_keys))
 
-    def _parse_column(self, column: str) -> tuple[str, str]:
-        """
-        Parses a column name into a tuple containing a data type and a key.
-
-        >>> r = Results.example()
-        >>> r._parse_column("answer.how_feeling")
-        ('answer', 'how_feeling')
-
-        The standard way a column is specified is with a dot-separated string, e.g. _parse_column("agent.status")
-        But you can also specify a single key, e.g. "status", in which case it will look up the data type.
-        """
-        if "." in column:
-            data_type, key = column.split(".")
-        else:
-            try:
-                data_type, key = self._key_to_data_type[column], column
-            except KeyError:
-                import difflib
-
-                close_matches = difflib.get_close_matches(
-                    column, self._key_to_data_type.keys()
-                )
-                if close_matches:
-                    suggestions = ", ".join(close_matches)
-                    raise ResultsColumnNotFoundError(
-                        f"Column '{column}' not found in data. Did you mean: {suggestions}?"
-                    )
-                else:
-                    raise ResultsColumnNotFoundError(
-                        f"Column {column} not found in data"
-                    )
-        return data_type, key
-
     def first(self) -> "Result":
         """Return the first observation in the results.
 
@@ -632,9 +677,11 @@ class Results(UserList, Mixins, Base):
         """
         if functions_dict is None:
             functions_dict = {}
-        return EvalWithCompoundTypes(
+        evaluator = EvalWithCompoundTypes(
             names=result.combined_dict, functions=functions_dict
         )
+        evaluator.functions.update(int=int, float=float)
+        return evaluator
 
     def mutate(
         self, new_var_string: str, functions_dict: Optional[dict] = None
@@ -721,8 +768,8 @@ class Results(UserList, Mixins, Base):
 
     def sample(
         self,
-        n: int = None,
-        frac: float = None,
+        n: Optional[int] = None,
+        frac: Optional[float] = None,
         with_replacement: bool = True,
         seed: Optional[str] = "edsl",
     ) -> Results:
@@ -771,13 +818,17 @@ class Results(UserList, Mixins, Base):
         Dataset([{'answer.how_feeling': ['OK', 'Great', 'Terrible', 'OK']}])
 
         >>> results.select('how_feeling', 'model', 'how_feeling')
-        Dataset([{'answer.how_feeling': ['OK', 'Great', 'Terrible', 'OK']}, {'model.model': ['gpt-4-1106-preview', 'gpt-4-1106-preview', 'gpt-4-1106-preview', 'gpt-4-1106-preview']}, {'answer.how_feeling': ['OK', 'Great', 'Terrible', 'OK']}])
+        Dataset([{'answer.how_feeling': ['OK', 'Great', 'Terrible', 'OK']}, {'model.model': ['...', '...', '...', '...']}, {'answer.how_feeling': ['OK', 'Great', 'Terrible', 'OK']}])
+
+        >>> from edsl import Results; r = Results.example(); r.select('answer.how_feeling_y')
+        Dataset([{'answer.how_feeling_yesterday': ['Great', 'Good', 'OK', 'Terrible']}])
         """
 
-        if len(self) == 0:
-            raise Exception("No data to select from---the Results object is empty.")
+        # if len(self) == 0:
+        #    raise Exception("No data to select from---the Results object is empty.")
 
         if not columns or columns == ("*",) or columns == (None,):
+            # is the users passes nothing, then we'll return all the columns
             columns = ("*.*",)
 
         if isinstance(columns[0], list):
@@ -801,6 +852,16 @@ class Results(UserList, Mixins, Base):
         # iterate through the passed columns
         for column in columns:
             # a user could pass 'result.how_feeling' or just 'how_feeling'
+            matches = self._matching_columns(column)
+            if len(matches) > 1:
+                raise Exception(
+                    f"Column '{column}' is ambiguous. Did you mean one of {matches}?"
+                )
+            if len(matches) == 0 and ".*" not in column:
+                raise Exception(f"Column '{column}' not found in data.")
+            if len(matches) == 1:
+                column = matches[0]
+
             parsed_data_type, parsed_key = self._parse_column(column)
             data_types = get_data_types_to_return(parsed_data_type)
             found_once = False  # we need to track this to make sure we found the key at least once
@@ -843,6 +904,21 @@ class Results(UserList, Mixins, Base):
 
         return Dataset(sorted_new_data)
 
+    def select(self, *columns: Union[str, list[str]]) -> "Results":
+        from edsl.results.Selector import Selector
+
+        if len(self) == 0:
+            raise Exception("No data to select from---the Results object is empty.")
+
+        selector = Selector(
+            known_data_types=self.known_data_types,
+            data_type_to_keys=self._data_type_to_keys,
+            key_to_data_type=self._key_to_data_type,
+            fetch_list_func=self._fetch_list,
+            columns=self.columns,
+        )
+        return selector.select(*columns)
+
     def sort_by(self, *columns: str, reverse: bool = False) -> Results:
         import warnings
 
@@ -851,6 +927,11 @@ class Results(UserList, Mixins, Base):
         )
         return self.order_by(*columns, reverse=reverse)
 
+    def _parse_column(self, column: str) -> tuple[str, str]:
+        if "." in column:
+            return column.split(".")
+        return self._key_to_data_type[column], column
+
     def order_by(self, *columns: str, reverse: bool = False) -> Results:
         """Sort the results by one or more columns.
 
@@ -948,7 +1029,9 @@ class Results(UserList, Mixins, Base):
         def has_single_equals(string):
             if "!=" in string:
                 return False
-            if "=" in string and not "==" in string:
+            if "=" in string and not (
+                "==" in string or "<=" in string or ">=" in string
+            ):
                 return True
 
         if has_single_equals(expression):
@@ -989,7 +1072,7 @@ class Results(UserList, Mixins, Base):
         return Results(survey=self.survey, data=new_data, created_columns=None)
 
     @classmethod
-    def example(cls, debug: bool = False, randomize: bool = False) -> Results:
+    def example(cls, randomize: bool = False) -> Results:
         """Return an example `Results` object.
 
         Example usage:
@@ -1003,7 +1086,13 @@ class Results(UserList, Mixins, Base):
 
         c = Cache()
         job = Jobs.example(randomize=randomize)
-        results = job.run(cache=c, debug=debug)
+        results = job.run(
+            cache=c,
+            stop_on_exception=True,
+            skip_retry=True,
+            raise_validation_errors=True,
+            disable_remote_inference=True,
+        )
         return results
 
     def rich_print(self):

edsl/results/ResultsDBMixin.py

@@ -136,9 +136,9 @@ class ResultsDBMixin:
 
         >>> from edsl.results import Results
         >>> r = Results.example()
-        >>> d = r.sql("select data_type, key, value from self where data_type = 'answer' limit 3", shape="long")
-        >>> list(d['value'])
-        ['OK', 'This is a real survey response from a human.', 'Great']
+        >>> d = r.sql("select data_type, key, value from self where data_type = 'answer' order by value limit 3", shape="long")
+        >>> sorted(list(d['value']))
+        ['Good', 'Great', 'Great']
 
         We can also return the data in wide format.
         Note the use of single quotes to escape the column names, as required by sql.

edsl/results/Selector.py

@@ -0,0 +1,118 @@
+from typing import Union, List, Dict, Any
+from collections import defaultdict
+from edsl.results.Dataset import Dataset
+
+
+class Selector:
+    def __init__(
+        self,
+        known_data_types: List[str],
+        data_type_to_keys: Dict[str, List[str]],
+        key_to_data_type: Dict[str, str],
+        fetch_list_func,
+        columns: List[str],
+    ):
+        self.known_data_types = known_data_types
+        self._data_type_to_keys = data_type_to_keys
+        self._key_to_data_type = key_to_data_type
+        self._fetch_list = fetch_list_func
+        self.columns = columns
+
+    def select(self, *columns: Union[str, List[str]]) -> "Dataset":
+        columns = self._normalize_columns(columns)
+        to_fetch = self._get_columns_to_fetch(columns)
+        new_data = self._fetch_data(to_fetch)
+        return Dataset(new_data)
+
+    def _normalize_columns(self, columns: Union[str, List[str]]) -> tuple:
+        if not columns or columns == ("*",) or columns == (None,):
+            return ("*.*",)
+        if isinstance(columns[0], list):
+            return tuple(columns[0])
+        return columns
+
+    def _get_columns_to_fetch(self, columns: tuple) -> Dict[str, List[str]]:
+        to_fetch = defaultdict(list)
+        self.items_in_order = []
+
+        for column in columns:
+            matches = self._find_matching_columns(column)
+            self._validate_matches(column, matches)
+
+            if len(matches) == 1:
+                column = matches[0]
+
+            data_type, key = self._parse_column(column)
+            self._process_column(data_type, key, to_fetch)
+
+        return to_fetch
+
+    def _find_matching_columns(self, partial_name: str) -> list[str]:
+        if "." in partial_name:
+            search_in_list = self.columns
+        else:
+            search_in_list = [s.split(".")[1] for s in self.columns]
+
+        matches = [s for s in search_in_list if s.startswith(partial_name)]
+        return [partial_name] if partial_name in matches else matches
+
+    def _validate_matches(self, column: str, matches: List[str]):
+        if len(matches) > 1:
+            raise ValueError(
+                f"Column '{column}' is ambiguous. Did you mean one of {matches}?"
+            )
+        if len(matches) == 0 and ".*" not in column:
+            raise ValueError(f"Column '{column}' not found in data.")
+
+    def _parse_column(self, column: str) -> tuple[str, str]:
+        if "." in column:
+            return column.split(".")
+        try:
+            return self._key_to_data_type[column], column
+        except KeyError:
+            self._raise_key_error(column)
+
+    def _raise_key_error(self, column: str):
+        import difflib
+
+        close_matches = difflib.get_close_matches(column, self._key_to_data_type.keys())
+        if close_matches:
+            suggestions = ", ".join(close_matches)
+            raise KeyError(
+                f"Column '{column}' not found in data. Did you mean: {suggestions}?"
+            )
+        else:
+            raise KeyError(f"Column {column} not found in data")
+
+    def _process_column(self, data_type: str, key: str, to_fetch: Dict[str, List[str]]):
+        data_types = self._get_data_types_to_return(data_type)
+        found_once = False
+
+        for dt in data_types:
+            relevant_keys = self._data_type_to_keys[dt]
+            for k in relevant_keys:
+                if k == key or key == "*":
+                    found_once = True
+                    to_fetch[dt].append(k)
+                    self.items_in_order.append(f"{dt}.{k}")
+
+        if not found_once:
+            raise ValueError(f"Key {key} not found in data.")
+
+    def _get_data_types_to_return(self, parsed_data_type: str) -> List[str]:
+        if parsed_data_type == "*":
+            return self.known_data_types
+        if parsed_data_type not in self.known_data_types:
+            raise ValueError(
+                f"Data type {parsed_data_type} not found in data. Did you mean one of {self.known_data_types}"
+            )
+        return [parsed_data_type]
+
+    def _fetch_data(self, to_fetch: Dict[str, List[str]]) -> List[Dict[str, Any]]:
+        new_data = []
+        for data_type, keys in to_fetch.items():
+            for key in keys:
+                entries = self._fetch_list(data_type, key)
+                new_data.append({f"{data_type}.{key}": entries})
+
+        return [d for key in self.items_in_order for d in new_data if key in d]