| import pandas as pd |
| import pytest |
|
|
| from expense_ml.data.sampling import balanced_training_sample |
|
|
|
|
| def test_balanced_training_sample_is_deterministic_and_bounded(): |
| frame = pd.DataFrame( |
| [ |
| {"text": f"usa-{i}", "label": "food", "source": "global", "country": "USA"} |
| for i in range(80) |
| ] |
| + [ |
| {"text": f"india-{i}", "label": "food", "source": "finee-india", "country": "India"} |
| for i in range(20) |
| ] |
| + [ |
| {"text": f"india-t-{i}", "label": "transportation", "source": "synthetic", "country": "India"} |
| for i in range(20) |
| ] |
| ) |
| a = balanced_training_sample(frame, max_rows=25, seed=17) |
| b = balanced_training_sample(frame, max_rows=25, seed=17) |
| assert len(a) == 25 |
| assert a["text"].tolist() == b["text"].tolist() |
| assert a["text"].nunique() == len(a) |
| assert set(a["country"]) == {"USA", "India"} |
| assert set(a["source"]) == {"global", "finee-india", "synthetic"} |
| assert set(a["label"]) == {"food", "transportation"} |
|
|
|
|
| def test_balanced_training_sample_rejects_impossible_label_budget(): |
| frame = pd.DataFrame( |
| [ |
| {"text": f"sample-{i}", "label": label, "source": "global", "country": "USA"} |
| for label in ("food", "transportation", "healthcare") |
| for i in range(10) |
| ] |
| ) |
| with pytest.raises(ValueError, match="too small to preserve all"): |
| balanced_training_sample(frame, max_rows=2, seed=1) |
|
|
