great-ai/great_ai/tracing/add_ground_truth.py

67 lines
1.9 KiB
Python

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[T](
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)