KPrashanth's picture
Deploy Rachana Data Studio
14954f2 verified
Raw
History Blame Contribute Delete
3.7 kB
from __future__ import annotations
import hashlib
import random
from collections import defaultdict
from datetime import datetime
from typing import Any, Callable
from pydantic import BaseModel, ConfigDict, Field
class SamplingManifestV1(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True)
schema_version: str = "1.0"
sample_id: str = Field(min_length=1)
strategy: str = Field(min_length=1)
population_hash: str = Field(min_length=1)
random_seed: str = Field(min_length=1)
selected_record_ids: tuple[str, ...]
strata_counts: dict[str, int]
def population_hash(record_ids: list[str]) -> str:
return hashlib.sha256("\n".join(sorted(record_ids)).encode("utf-8")).hexdigest()
def reproducible_random_sample(record_ids: list[str], size: int, seed: str) -> SamplingManifestV1:
unique = sorted(set(record_ids))
if size < 1 or size > len(unique):
raise ValueError("Sample size must be within the unique population size.")
selected = sorted(random.Random(seed).sample(unique, size))
return SamplingManifestV1(
sample_id=hashlib.sha256(f"random:{seed}:{population_hash(unique)}:{size}".encode("utf-8")).hexdigest(),
strategy="random",
population_hash=population_hash(unique),
random_seed=seed,
selected_record_ids=tuple(selected),
strata_counts={"all": len(selected)},
)
def stratified_sample(
records: list[dict[str, Any]],
size: int,
seed: str,
stratum: Callable[[dict[str, Any]], str],
) -> SamplingManifestV1:
groups: dict[str, list[str]] = defaultdict(list)
for record in records:
groups[stratum(record)].append(record["record_id"])
if size < len(groups):
raise ValueError("Sample size must include at least one record per stratum.")
population = sorted({record_id for values in groups.values() for record_id in values})
if size > len(population):
raise ValueError("Sample size exceeds population.")
rng = random.Random(seed)
selected: list[str] = []
counts: dict[str, int] = {}
remaining = size
strata = sorted(groups)
for index, key in enumerate(strata):
available = sorted(set(groups[key]))
target = max(1, round(size * len(available) / len(population)))
target = min(target, len(available), remaining - (len(strata) - index - 1))
picked = rng.sample(available, target)
selected.extend(picked)
counts[key] = len(picked)
remaining -= len(picked)
if remaining:
candidates = sorted(set(population).difference(selected))
selected.extend(rng.sample(candidates, remaining))
return SamplingManifestV1(
sample_id=hashlib.sha256(f"stratified:{seed}:{population_hash(population)}:{size}".encode("utf-8")).hexdigest(),
strategy="stratified",
population_hash=population_hash(population),
random_seed=seed,
selected_record_ids=tuple(sorted(selected)),
strata_counts=counts,
)
def capacity_report(review_events: list[dict[str, Any]]) -> dict[str, Any]:
finalized = [event for event in review_events if event.get("outcome") not in {"skip", None}]
hours = sum(float(event.get("duration_seconds", 0)) for event in finalized) / 3600
disagreements = sum(bool(event.get("disagreed")) for event in finalized)
return {
"finalized_reviews": len(finalized),
"review_hours": round(hours, 2),
"reviews_per_hour": round(len(finalized) / hours, 2) if hours else None,
"disagreement_rate": round(disagreements / len(finalized), 4) if finalized else None,
"generated_at": datetime.utcnow().isoformat() + "Z",
}