Spaces:
Restarting
Restarting
| import pandas as pd | |
| from expense_ml.data.split import remove_conflicting_text_groups, split_dataset | |
| def test_split_is_deterministic_and_has_no_exact_text_leakage(): | |
| rows = [] | |
| for i in range(50): | |
| rows.append({"text": f"merchant {i}", "label": "food", "source": "global", "country": "USA"}) | |
| rows.append({"text": f"merchant {i + 50}", "label": "transportation", "source": "global", "country": "Canada"}) | |
| rows.extend( | |
| [ | |
| {"text": "repeated merchant", "label": "food", "source": "global", "country": "USA"}, | |
| {"text": "repeated merchant", "label": "food", "source": "finee-india", "country": "India"}, | |
| {"text": "repeated merchant", "label": "food", "source": "synthetic-indian-transactions", "country": "India"}, | |
| ] | |
| ) | |
| frame = pd.DataFrame(rows) | |
| a = split_dataset(frame, seed=7, test_size=.2, validation_size=.1) | |
| b = split_dataset(frame, seed=7, test_size=.2, validation_size=.1) | |
| assert a.train["text"].tolist() == b.train["text"].tolist() | |
| for left, right in ( | |
| (a.train, a.test), | |
| (a.train, a.validation), | |
| (a.train, a.india_holdout), | |
| (a.validation, a.test), | |
| (a.validation, a.india_holdout), | |
| (a.test, a.india_holdout), | |
| ): | |
| assert set(left["text"]) & set(right["text"]) == set() | |
| combined_partition_texts = set().union(*(set(part["text"]) for part in (a.train, a.validation, a.test, a.india_holdout))) | |
| assert "repeated merchant" in combined_partition_texts | |
| assert "repeated merchant" not in set(a.india_holdout["text"]) | |
| def test_cross_country_duplicate_is_excluded_from_india_holdout(): | |
| frame = pd.DataFrame( | |
| [ | |
| {"text": "same narration", "label": "food", "source": "global", "country": "USA"}, | |
| {"text": "same narration", "label": "food", "source": "global", "country": "India"}, | |
| ] | |
| + [ | |
| {"text": f"india {i}", "label": "food", "source": "global", "country": "India"} | |
| for i in range(100) | |
| ] | |
| + [ | |
| {"text": f"usa {i}", "label": "transportation", "source": "global", "country": "USA"} | |
| for i in range(100) | |
| ] | |
| ) | |
| result = split_dataset(frame, seed=17, test_size=.2, validation_size=.1) | |
| assert "same narration" not in set(result.india_holdout["text"]) | |
| def test_split_uses_country_metadata_for_india_holdout(): | |
| frame = pd.DataFrame( | |
| [ | |
| {"text": f"india {i}", "label": "food", "source": "global", "country": "India"} | |
| for i in range(100) | |
| ] | |
| + [ | |
| {"text": f"usa {i}", "label": "transportation", "source": "global", "country": "USA"} | |
| for i in range(100) | |
| ] | |
| ) | |
| result = split_dataset(frame, seed=11, test_size=.2, validation_size=.1) | |
| assert len(result.india_holdout) > 0 | |
| assert set(result.india_holdout["country"]) == {"India"} | |
| def test_conflicting_text_groups_are_removed_before_splitting(): | |
| frame = pd.DataFrame( | |
| [ | |
| {"text": "ambiguous narration", "label": "food_dining", "source": "global", "country": "USA"}, | |
| {"text": "ambiguous narration", "label": "shopping_retail", "source": "finee", "country": "India"}, | |
| {"text": "clear food", "label": "food_dining", "source": "global", "country": "USA"}, | |
| {"text": "clear transport", "label": "transportation", "source": "global", "country": "USA"}, | |
| ] | |
| ) | |
| cleaned, summary = remove_conflicting_text_groups(frame) | |
| assert "ambiguous narration" not in set(cleaned["text"]) | |
| assert set(cleaned["text"]) == {"clear food", "clear transport"} | |
| assert summary["conflicting_text_groups"] == 1 | |
| assert summary["rows_removed"] == 2 | |
| def test_split_removes_conflicting_groups_before_partitioning(): | |
| rows = [ | |
| {"text": "ambiguous narration", "label": "food_dining", "source": "global", "country": "USA"}, | |
| {"text": "ambiguous narration", "label": "shopping_retail", "source": "finee", "country": "India"}, | |
| ] | |
| rows.extend( | |
| {"text": f"food {i}", "label": "food_dining", "source": "global", "country": "USA"} | |
| for i in range(30) | |
| ) | |
| rows.extend( | |
| {"text": f"transport {i}", "label": "transportation", "source": "global", "country": "USA"} | |
| for i in range(30) | |
| ) | |
| result = split_dataset(pd.DataFrame(rows), seed=19, test_size=.2, validation_size=.1) | |
| partition_texts = set().union( | |
| *(set(part["text"]) for part in (result.train, result.validation, result.test, result.india_holdout)) | |
| ) | |
| assert "ambiguous narration" not in partition_texts | |