Move files

great_ai/tracing/add_ground_truth.py

@@ -0,0 +1,67 @@
+from datetime import datetime
+from math import ceil
+from random import shuffle
+from typing import Any, Iterable, List, TypeVar
+
+from ..constants import (
+    GROUND_TRUTH_TAG_NAME,
+    TEST_SPLIT_TAG_NAME,
+    TRAIN_SPLIT_TAG_NAME,
+    VALIDATION_SPLIT_TAG_NAME,
+)
+from ..context import get_context
+from ..views import Trace
+
+T = TypeVar("T")
+
+
+def add_ground_truth(
+    inputs: Iterable[Any],
+    expected_outputs: Iterable[T],
+    *,
+    tags: List[str] = [],
+    train_split_ratio: float = 1,
+    test_split_ratio: float = 0,
+    validation_split_ratio: float = 0
+) -> None:
+    get_context()  # this resets the seed
+
+    inputs = list(inputs)
+    expected_outputs = list(expected_outputs)
+    assert len(inputs) == len(
+        expected_outputs
+    ), "The length of the inputs and expected_outputs must be equal"
+
+    sum_ratio = train_split_ratio + test_split_ratio + validation_split_ratio
+    assert sum_ratio > 0, "The sum of the split ratios must be a positive number"
+
+    train_split_ratio /= sum_ratio
+    test_split_ratio /= sum_ratio
+    validation_split_ratio /= sum_ratio
+
+    values = list(zip(inputs, expected_outputs))
+    shuffle(values)
+
+    split_tags = (
+        [TRAIN_SPLIT_TAG_NAME] * ceil(train_split_ratio * len(inputs))
+        + [TEST_SPLIT_TAG_NAME] * ceil(test_split_ratio * len(inputs))
+        + [VALIDATION_SPLIT_TAG_NAME] * ceil(validation_split_ratio * len(inputs))
+    )
+    shuffle(split_tags)
+
+    created = datetime.utcnow().isoformat()
+    traces = [
+        Trace(
+            created=created,
+            original_execution_time_ms=0,
+            logged_values=X if isinstance(X, dict) else {"input": X},
+            models=[],
+            output=y,
+            feedback=y,
+            exception=None,
+            tags=[GROUND_TRUTH_TAG_NAME, split_tag, *tags],
+        )
+        for ((X, y), split_tag) in zip(values, split_tags)
+    ]
+
+    get_context().tracing_database.save_batch(traces)