diff --git a/fern/versions/latest/pages/concepts/person_sampling.mdx b/fern/versions/latest/pages/concepts/person_sampling.mdx index c6ee76f85..0f7240e1b 100644 --- a/fern/versions/latest/pages/concepts/person_sampling.mdx +++ b/fern/versions/latest/pages/concepts/person_sampling.mdx @@ -61,12 +61,14 @@ Supported locales: - `en_US`: United States - `en_IN`: India (English) - `en_SG`: Singapore (English) +- `es_SV`: El Salvador (Spanish) - `fr_FR`: France (French) - `hi_Deva_IN`: India (Devanagari script) - `hi_Latn_IN`: India (Latin script) - `ja_JP`: Japan - `ko_KR`: South Korea (Korean) - `pt_BR`: Brazil (Portuguese) +- `vi_VN`: Vietnam (Vietnamese) ### Features - **Demographically accurate personal details**: Names, ages, sex, marital status, education, occupation based on census data @@ -268,6 +270,17 @@ The managed reader normalizes `Weiblich` / `Männlich` to `Female` / `Male`, so - `race` - Census-reported race +**El Salvador-Specific Fields (`es_SV`):** + +- `care_access` - Access to healthcare +- `disability_status` - Disability status +- `income_band` - Income range +- `languages_spoken` - Languages spoken + +**Vietnam-Specific Fields (`vi_VN`):** + +- `household_size` - Number of household members + **Singapore-Specific Fields (`en_SG`):** - `industry` - Industry of employment @@ -321,7 +334,7 @@ The managed reader normalizes `Weiblich` / `Männlich` to `Female` / `Male`, so | Parameter | Type | Description | |-----------|------|-------------| -| `locale` | str | Language/region code - must be one of: "en_GB", "de_DE", "en_US", "en_IN", "en_SG", "fr_FR", "hi_Deva_IN", "hi_Latn_IN", "ja_JP", "ko_KR", "pt_BR" | +| `locale` | str | Language/region code - must be one of: "en_GB", "de_DE", "en_US", "en_IN", "en_SG", "es_SV", "fr_FR", "hi_Deva_IN", "hi_Latn_IN", "ja_JP", "ko_KR", "pt_BR", "vi_VN" | | `sex` | str (optional) | Filter by "Male" or "Female" | | `city` | str or list[str] (optional) | Filter by specific city or cities within locale | | `age_range` | list[int] (optional) | Two-element list [min_age, max_age] (default: [18, 114]) | diff --git a/packages/data-designer-config/src/data_designer/config/utils/constants.py b/packages/data-designer-config/src/data_designer/config/utils/constants.py index 6ea435b5f..30452e760 100644 --- a/packages/data-designer-config/src/data_designer/config/utils/constants.py +++ b/packages/data-designer-config/src/data_designer/config/utils/constants.py @@ -399,12 +399,14 @@ class NordColor(Enum): "en_US": "1.24 GB", "en_IN": "2.39 GB", "en_SG": "0.30 GB", + "es_SV": "0.74 GB", "fr_FR": "3.87 GB", "hi_Deva_IN": "4.14 GB", "hi_Latn_IN": "2.7 GB", "ja_JP": "1.69 GB", "ko_KR": "2.66 GB", "pt_BR": "2.33 GB", + "vi_VN": "2.58 GB", } LOCALES_WITH_MANAGED_DATASETS = list[str](NEMOTRON_PERSONAS_DATASET_SIZES.keys()) diff --git a/packages/data-designer-engine/src/data_designer/engine/sampling_gen/entities/dataset_based_person_fields.py b/packages/data-designer-engine/src/data_designer/engine/sampling_gen/entities/dataset_based_person_fields.py index 711e06c38..8e49ff48d 100644 --- a/packages/data-designer-engine/src/data_designer/engine/sampling_gen/entities/dataset_based_person_fields.py +++ b/packages/data-designer-engine/src/data_designer/engine/sampling_gen/entities/dataset_based_person_fields.py @@ -83,6 +83,13 @@ "nationality", # pt_BR-specific "race", + # es_SV-specific + "care_access", + "disability_status", + "income_band", + "languages_spoken", + # vi_VN-specific + "household_size", # India locales (en_IN, hi_Deva_IN, hi_Latn_IN) "education_degree", "first_language", diff --git a/packages/data-designer-engine/tests/engine/sampling_gen/test_people_gen.py b/packages/data-designer-engine/tests/engine/sampling_gen/test_people_gen.py index 310c155a9..26a78efc8 100644 --- a/packages/data-designer-engine/tests/engine/sampling_gen/test_people_gen.py +++ b/packages/data-designer-engine/tests/engine/sampling_gen/test_people_gen.py @@ -153,3 +153,48 @@ def test_managed_uk_germany_sampling( assert (field in person) == with_synthetic_personas if with_synthetic_personas: assert person[field] == record[field] + + +@pytest.mark.parametrize( + "locale,field,selected_value,other_value", + [ + ("es_SV", "care_access", "Public", "Private"), + ("es_SV", "disability_status", "None", "Physical"), + ("es_SV", "income_band", "Low", "High"), + ("es_SV", "languages_spoken", "Spanish", "English"), + ("vi_VN", "household_size", 4, 2), + ], +) +@pytest.mark.parametrize("with_synthetic_personas", [False, True]) +def test_managed_locale_demographic_output_and_filtering( + tmp_path: Path, + locale: str, + field: str, + selected_value: str | int, + other_value: str | int, + with_synthetic_personas: bool, +) -> None: + dataset_dir = tmp_path / "datasets" + dataset_dir.mkdir() + record = { + "first_name": "Alex", + "last_name": "Smith", + "age": 35, + "persona": "Enjoys learning languages.", + field: selected_value, + } + lazy.pd.DataFrame([record, {**record, field: other_value}]).to_parquet(dataset_dir / f"{locale}.parquet") + generator = PeopleGenFromDataset(load_person_data_sampler(LocalPersonReader(tmp_path), locale), locale) + params = PersonSamplerParams( + locale=locale, + select_field_values={field: [str(selected_value)]}, + with_synthetic_personas=with_synthetic_personas, + ) + sampler = PersonSampler(params, people_gen_resource={params.people_gen_key: generator}) + + people = sampler.sample(1) + + assert len(people) == 1 + assert all(person[field] == selected_value for person in people) + assert all(person["locale"] == locale for person in people) + assert all(("persona" in person) == with_synthetic_personas for person in people) diff --git a/packages/data-designer/tests/cli/controllers/test_download_controller.py b/packages/data-designer/tests/cli/controllers/test_download_controller.py index b449132b0..930ccb283 100644 --- a/packages/data-designer/tests/cli/controllers/test_download_controller.py +++ b/packages/data-designer/tests/cli/controllers/test_download_controller.py @@ -98,20 +98,10 @@ def test_run_personas_with_all_flag( # Verify NGC check was called mock_check_ngc.assert_called_once() - # Verify all locales were downloaded - assert mock_download.call_count == 11 - - # Verify each locale was downloaded + expected_locales = {locale.code for locale in controller.persona_repository.list_all()} downloaded_locales = [call[0][0] for call in mock_download.call_args_list] - assert "en_US" in downloaded_locales - assert "en_IN" in downloaded_locales - assert "en_SG" in downloaded_locales - assert "fr_FR" in downloaded_locales - assert "hi_Deva_IN" in downloaded_locales - assert "hi_Latn_IN" in downloaded_locales - assert "ja_JP" in downloaded_locales - assert "ko_KR" in downloaded_locales - assert "pt_BR" in downloaded_locales + assert set(downloaded_locales) == expected_locales + assert len(downloaded_locales) == len(expected_locales) @patch.object(DownloadController, "_download_locale", return_value=True) @@ -260,16 +250,9 @@ def test_determine_locales_with_all_flag(controller: DownloadController) -> None """Test _determine_locales returns all locales when all_locales=True.""" result = controller._determine_locales(locales=None, all_locales=True) - assert len(result) == 11 - assert "en_US" in result - assert "en_IN" in result - assert "en_SG" in result - assert "fr_FR" in result - assert "hi_Deva_IN" in result - assert "hi_Latn_IN" in result - assert "ja_JP" in result - assert "ko_KR" in result - assert "pt_BR" in result + expected_locales = {locale.code for locale in controller.persona_repository.list_all()} + assert set(result) == expected_locales + assert len(result) == len(expected_locales) def test_determine_locales_with_valid_locale_flags(controller: DownloadController) -> None: diff --git a/packages/data-designer/tests/cli/repositories/test_persona_repository.py b/packages/data-designer/tests/cli/repositories/test_persona_repository.py index 7c205d948..0ee1f6946 100644 --- a/packages/data-designer/tests/cli/repositories/test_persona_repository.py +++ b/packages/data-designer/tests/cli/repositories/test_persona_repository.py @@ -17,7 +17,6 @@ def repository() -> PersonaRepository: def test_init(repository: PersonaRepository) -> None: """Test repository initialization creates registry.""" assert repository._registry is not None - assert len(repository._registry.locales) == 11 assert repository._registry.dataset_prefix == "nemotron-personas-dataset-" @@ -26,24 +25,6 @@ def test_list_all(repository: PersonaRepository) -> None: locales = repository.list_all() assert isinstance(locales, list) - assert len(locales) == 11 - - # Verify all expected locales are present - locale_codes = {locale.code for locale in locales} - assert locale_codes == { - "en_GB", - "de_DE", - "en_US", - "en_IN", - "en_SG", - "fr_FR", - "hi_Deva_IN", - "hi_Latn_IN", - "ja_JP", - "ko_KR", - "pt_BR", - } - # Verify each locale has required fields for locale in locales: assert isinstance(locale, PersonaLocale) @@ -65,16 +46,23 @@ def test_get_by_code_valid_locale(repository: PersonaRepository) -> None: def test_get_by_code_all_locales(repository: PersonaRepository) -> None: """Test getting each locale by code.""" test_cases = [ + ("es_SV", "0.74 GB", "nemotron-personas-dataset-es_sv"), + ("vi_VN", "2.58 GB", "nemotron-personas-dataset-vi_vn"), ("en_GB", "2.03 GB", "nemotron-personas-dataset-en_gb"), ("de_DE", "2.55 GB", "nemotron-personas-dataset-de_de"), ("en_US", "1.24 GB", "nemotron-personas-dataset-en_us"), ("en_IN", "2.39 GB", "nemotron-personas-dataset-en_in"), + ("en_SG", "0.30 GB", "nemotron-personas-dataset-en_sg"), + ("fr_FR", "3.87 GB", "nemotron-personas-dataset-fr_fr"), + ("pt_BR", "2.33 GB", "nemotron-personas-dataset-pt_br"), ("hi_Deva_IN", "4.14 GB", "nemotron-personas-dataset-hi_deva_in"), ("hi_Latn_IN", "2.7 GB", "nemotron-personas-dataset-hi_latn_in"), ("ja_JP", "1.69 GB", "nemotron-personas-dataset-ja_jp"), ("ko_KR", "2.66 GB", "nemotron-personas-dataset-ko_kr"), ] + assert {locale.code for locale in repository.list_all()} == {code for code, _, _ in test_cases} + for code, expected_size, expected_dataset_name in test_cases: locale = repository.get_by_code(code) assert locale is not None diff --git a/packages/data-designer/tests/cli/services/test_download_service.py b/packages/data-designer/tests/cli/services/test_download_service.py index 9670afc78..ae0bc5ed6 100644 --- a/packages/data-designer/tests/cli/services/test_download_service.py +++ b/packages/data-designer/tests/cli/services/test_download_service.py @@ -67,7 +67,7 @@ def test_get_available_locales(service: DownloadService) -> None: locales = service.get_available_locales() assert isinstance(locales, dict) - assert len(locales) == 11 + assert set(locales) == {locale.code for locale in service.persona_repository.list_all()} assert "en_GB" in locales assert "de_DE" in locales assert "en_US" in locales