great-ai/great_ai/tracing/add_ground_truth.py
2022-07-08 20:33:33 +02:00

72 lines
2.1 KiB
Python

from datetime import datetime
from math import ceil
from random import shuffle
from typing import Any, Iterable, List, TypeVar, cast
from uuid import uuid4
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 = [
cast(
Trace[T],
Trace( # avoid ValueError: "Trace" object has no field "__orig_class__"
trace_id=str(uuid4()),
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)