Add docstrings

great_ai/tracing/add_ground_truth.py

@@ -25,7 +25,37 @@ def add_ground_truth(
     test_split_ratio: float = 0,
     validation_split_ratio: float = 0
 ) -> None:
-    get_context()  # this resets the seed
+    """Add training data (with optional train-test splitting).
+
+    Add and tag datapoints, wrap them into traces. The `inputs` are available via the
+    `.input` property, while `expected_outputs` under both the `.output` and `.feedback`
+    properties.
+
+    All generated traces are tagged with `ground_truth` by default. Additional tags can
+    be also provided. Using the `split_ratio` arguments, tags can be given randomly with
+    a user-defined distribution. Only the ratio of the splits matter, they don't have to
+    add up to 1.
+
+    Examples:
+        >>> add_ground_truth(
+        ...    [1, 2, 3],
+        ...    ['odd', 'even', 'odd'],
+        ...    tags='my_tag',
+        ...    train_split_ratio=1,
+        ...    test_split_ratio=1,
+        ...    validation_split_ratio=0.5,
+        ... )
+
+    Args:
+        inputs: The inputs. (X in scikit-learn)
+        expected_outputs: The ground-truth values corresponding to the inputs. (y in
+            scikit-learn)
+        tags: A single tag or a list of tags to append to each generated trace's tags.
+        train_split_ratio: The probability-weight of giving each trace the `train` tag.
+        test_split_ratio: The probability-weight of giving each trace the `test` tag.
+        validation_split_ratio: The probability-weight of giving each trace the
+        `validation` tag.
+    """
 
     inputs = list(inputs)
     expected_outputs = list(expected_outputs)

great_ai/tracing/delete_ground_truth.py

@@ -7,9 +7,28 @@ from ..context import get_context
 def delete_ground_truth(
     conjunctive_tags: Union[List[str], str] = [],
     *,
-    until: Optional[datetime] = None,
     since: Optional[datetime] = None,
+    until: Optional[datetime] = None,
 ) -> None:
+    """Delete traces matching the given criteria.
+
+    Takes the same arguments as `query_ground_truth` but instead of returning them,
+    it simply deletes them.
+
+    You can rely on the efficiency of the delete's implementation.
+
+    Examples:
+        >>> delete_ground_truth(['train', 'test', 'validation'])
+
+    Args:
+        conjunctive_tags: Single tag or a list of tags which the deleted traces have to
+            match. The relationship between the tags is conjunctive (AND).
+        since: Only delete traces created after the given timestamp. `None` means no
+            filtering.
+        until: Only delete traces created before the given timestamp. `None` means no
+            filtering.
+    """
+
     tags = (
         conjunctive_tags if isinstance(conjunctive_tags, list) else [conjunctive_tags]
     )

great_ai/tracing/query_ground_truth.py

@@ -9,14 +9,43 @@ def query_ground_truth(
     conjunctive_tags: Union[List[str], str] = [],
     *,
     since: Optional[datetime] = None,
+    until: Optional[datetime] = None,
     return_max_count: Optional[int] = None
 ) -> List[Trace]:
+    """Return training samples.
+
+    Combines, filters, and returns datapoints that have been either added by
+    `add_ground_truth` or were the result of a prediction after which their trace got
+    a feedback through the RESP API-s `/traces/{trace_id}/feedback` endpoint
+    (end-to-end feedback).
+
+    Filtering can be used to only return points matching all given tags (or the single
+    given tag) and by time of creation.
+
+    Examples:
+        >>> query_ground_truth()
+        [...]
+
+    Args:
+        conjunctive_tags: Single tag or a list of tags which the returned traces have to
+            match. The relationship between the tags is conjunctive (AND).
+        since: Only return traces created after the given timestamp. `None` means no
+            filtering.
+        until: Only return traces created before the given timestamp. `None` means no
+            filtering.
+        return_max_count: Return at-most this many traces. (take, limit)
+    """
+
     tags = (
         conjunctive_tags if isinstance(conjunctive_tags, list) else [conjunctive_tags]
     )
     db = get_context().tracing_database
 
     items, length = db.query(
-        conjunctive_tags=tags, since=since, take=return_max_count, has_feedback=True
+        conjunctive_tags=tags,
+        since=since,
+        until=until,
+        take=return_max_count,
+        has_feedback=True,
     )
     return items

great_ai/tracing/tracing_context.py

@@ -13,6 +13,11 @@ T = TypeVar("T")
 
 
 class TracingContext(Generic[T]):
+    """Provide a global context variable throughout the life-cycle of a prediction.
+
+    Should only be used by internal great-ai functions.
+    """
+
     def __init__(self, function_name: str, do_not_persist_traces: bool) -> None:
         self._do_not_persist_traces = do_not_persist_traces
         self._models: List[Model] = []
@@ -34,9 +39,10 @@ class TracingContext(Generic[T]):
         assert self._trace is None, "has been already finalised"
 
         delta_time = round((perf_counter() - self._start_time) * 1000, 4)
-        self._trace = cast(
+        self._trace = (
+            cast(  # avoid ValueError: "Trace" object has no field "__orig_class__"
                 Trace[T],
-            Trace(  # avoid ValueError: "Trace" object has no field "__orig_class__"
+                Trace(
                     trace_id=str(uuid4()),
                     created=self._start_datetime.isoformat(),
                     original_execution_time_ms=delta_time,
@@ -55,6 +61,7 @@ class TracingContext(Generic[T]):
                     ],
                 ),
             )
+        )
 
         return self._trace
 

great_ai/utilities/unchunk.py

@@ -9,6 +9,8 @@ def unchunk(chunks: Iterable[Optional[Iterable[T]]]) -> Iterable[T]:
     The inverse operation of [chunk][great_ai.utilities.chunk.chunk].
     Useful for parallel processing.
 
+    Similar to itertools.chain but ignores `None` chunks.
+
     Examples:
         >>> list(unchunk([[0, 1, 2], [3, 4, 5], [6, 7, 8], [9]]))
         [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]