diff --git a/dpsynth/domain.py b/dpsynth/domain.py index 86c19c4..6275787 100644 --- a/dpsynth/domain.py +++ b/dpsynth/domain.py @@ -56,7 +56,8 @@ PathType = pathlib.Path -CategoricalValue: TypeAlias = bool | int | str +CategoricalValue: TypeAlias = bool | int | float | str + IntervalHandling = Literal['midpoint', 'sample', 'interval'] diff --git a/dpsynth/local_mode/vectorized_transformations.py b/dpsynth/local_mode/vectorized_transformations.py index 401dc41..4e1b0dd 100644 --- a/dpsynth/local_mode/vectorized_transformations.py +++ b/dpsynth/local_mode/vectorized_transformations.py @@ -179,7 +179,7 @@ def undiscretize( ``attribute_domain.resolved_sentinel``. """ rng = np.random.default_rng(rng) - min_, max_ = attribute_domain.min_value, attribute_domain.max_value + min_, max_ = attribute_domain.exclusive_min_value, attribute_domain.max_value _validate_bin_edges(bin_edges, attribute_domain) if bin_edges.size == 0: @@ -195,8 +195,10 @@ def undiscretize( sentinel = attribute_domain.resolved_sentinel if handling == 'interval': - # First interval is closed on both sides; the rest are half-open. - strs = [f'[{lefts[0]}, {rights[0]}]'] + [ + # First interval is closed on both sides (display uses min_value for + # readability); the rest are half-open (left, right]. + display_min = attribute_domain.min_value + strs = [f'[{display_min}, {rights[0]}]'] + [ f'({l}, {r}]' for l, r in zip(lefts[1:], rights[1:]) ] values = np.array(strs, dtype=str) diff --git a/dpsynth/text/bulk_inference.py b/dpsynth/text/bulk_inference.py deleted file mode 100644 index 077efdb..0000000 --- a/dpsynth/text/bulk_inference.py +++ /dev/null @@ -1,343 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Bulk LLM inference for synthetic text generation.""" - -from collections.abc import Callable, Sequence -import concurrent.futures -import dataclasses -import enum -import functools -import random -import re -import time -from typing import Protocol, TypeVar - -from absl import logging -from google import genai -from google.genai import types -import pandas as pd -import pydantic - - -class ModelName(enum.StrEnum): - """Model names supported by the google.genai API.""" - - GEMINI_2_5_FLASH_LITE = 'gemini-2.5-flash-lite' - GEMINI_3_FLASH = 'gemini-3-flash-preview' - GEMINI_3_5_FLASH = 'gemini-3.5-flash' - GEMMA_4_27B = 'gemma-4-26b-a4b-it' - GEMMA_4_31B = 'gemma-4-31b-it' - - -class TextGenerationBackend(Protocol): - """Interface for bulk LLM inference operations. - - Implementations provide the two LLM inference capabilities needed by the - synthetic text generation pipeline: - - 1. **Annotation**: extracting structured categorical features from text, - typically via constrained decoding with a pydantic response schema. - 2. **Generation**: producing free-form synthetic text conditioned on features. - - **Index-alignment guarantee**: both methods return output that is - positionally aligned with the input sequence — ``len(output) == len(input)`` - always holds. Failed items are represented as null rows (annotation) or - empty strings (generation). - """ - - def annotate( - self, - texts: Sequence[str], - schema: type[pydantic.BaseModel], - system_prompt: str, - ) -> pd.DataFrame: - """Extract structured features from texts using an LLM. - - Args: - texts: Input texts to annotate. - schema: Pydantic model class defining the features to extract. Fields may - use ``Literal`` type annotations for constrained decoding or plain types - such as ``str`` for open-ended annotation. Field names and descriptions - guide the LLM. - system_prompt: System-level instructions for the LLM describing how to - annotate the texts. - - Returns: - A DataFrame with exactly ``len(texts)`` rows and one column per field - in ``schema``. Rows where annotation failed contain ``None`` values. - """ - ... - - def generate(self, prompts: Sequence[str]) -> list[str]: - """Generate free-form text from prompts. - - Args: - prompts: Fully constructed prompts, each describing the desired output - including the target features and any formatting requirements. - - Returns: - A list of exactly ``len(prompts)`` strings. Failed generations are - represented as empty strings. - """ - ... - - -@dataclasses.dataclass(frozen=True) -class GenAIBackend: - """TextGenerationBackend using the google.genai API. - - Uses ``client.models.generate_content()`` for both annotation (with - structured output via ``response_schema``) and free-form generation. - - Attributes: - model: Model name string (e.g., ``'gemini-2.5-flash-lite'``). Accepts any - ``ModelName`` enum value or arbitrary string for unlisted models. - api_key: API key for authentication. - poll_interval_seconds: How often to poll for batch job completion. - chunk_size: Number of texts per batch job. - max_concurrent_jobs: Maximum number of active parallel batch jobs. - """ - - model: str = ModelName.GEMINI_2_5_FLASH_LITE - api_key: str | None = None - poll_interval_seconds: int = 60 - chunk_size: int = 100 - max_concurrent_jobs: int = 8 - - @functools.cached_property - def client(self) -> genai.Client: - """Creates and caches a genai.Client.""" - kwargs = {'http_options': types.HttpOptions(api_version='v1alpha')} - - if self.api_key: - kwargs['api_key'] = self.api_key # pyrefly: ignore[bad-assignment] - return genai.Client(**kwargs) - - def _parse_job_responses( - self, - batch_job: types.BatchJob, - schema: type[pydantic.BaseModel], - ) -> list[dict[str, str | None]]: - """Parses responses from a completed BatchJob.""" - if batch_job.state != types.JobState.JOB_STATE_SUCCEEDED: - error_msg = ( - f'Batch job {batch_job.name} ended with state={batch_job.state}.' - ) - if batch_job.error: - error_msg += f' Error: {batch_job.error}' - raise RuntimeError(error_msg) - - null_row = {f: None for f in schema.model_fields.keys()} - - inlined_responses = ( - batch_job.dest.inlined_responses if batch_job.dest else [] - ) or [] - - chunk_rows = [] - for i, inlined_resp in enumerate(inlined_responses): - row = dict(null_row) # Default to null row - try: - if inlined_resp.error: - logging.warning( - 'Batch result %d in job %s had error: %s', - i, - batch_job.name, - inlined_resp.error, - ) - else: - response_text = inlined_resp.response and inlined_resp.response.text - if response_text: - row = schema.model_validate_json( - _strip_markdown_fences(response_text) - ).model_dump() - else: - logging.warning( - 'Empty batch response in job %s for text %d.', - batch_job.name, - i, - ) - except Exception as e: # pylint: disable=broad-except - logging.warning( - 'Failed to parse batch result %d in job %s: %s', - i, - batch_job.name, - e, - ) - chunk_rows.append(row) - return chunk_rows - - def _submit_and_poll_chunk( - self, - chunk_texts: Sequence[str], - config: types.GenerateContentConfig | None = None, - ) -> types.BatchJob: - """Submit a batch job for one chunk and poll until done.""" - - inlined_requests = [ - types.InlinedRequest(contents=text, config=config) - for text in chunk_texts - ] - - job = _call_with_retry( - lambda: self.client.batches.create( - model=self.model, src=inlined_requests - ), - 'create', - ) - logging.info('Batch annotate: job %s created.', job.name) - - while not job.done: - time.sleep(self.poll_interval_seconds) - job = _call_with_retry( - lambda: self.client.batches.get(name=job.name), 'get' # pyrefly: ignore[bad-argument-type] - ) - - logging.info( - 'Batch annotate: job %s completed with state=%s', - job.name, - job.state, - ) - return job - - def annotate( - self, - texts: Sequence[str], - schema: type[pydantic.BaseModel], - system_prompt: str, - ) -> pd.DataFrame: - """Extract structured features via the GenAI Batch API. - - Submits texts as inlined requests to the batch prediction endpoint, - polls for completion, and parses the inlined responses. - - Args: - texts: Input texts to annotate. - schema: Pydantic model used as the ``response_schema``. - system_prompt: System-level instructions for the LLM. - - Returns: - DataFrame with exactly ``len(texts)`` rows. Failed rows have ``None``. - - Raises: - RuntimeError: If the batch job fails or is cancelled. - """ - config = types.GenerateContentConfig( - system_instruction=system_prompt, - response_mime_type='application/json', - response_schema=schema, - ) - - chunks = [ - texts[i : i + self.chunk_size] - for i in range(0, len(texts), self.chunk_size) - ] - - logging.info( - 'Batch annotate: processing %d chunks with concurrency limit %d...', - len(chunks), - self.max_concurrent_jobs, - ) - - with concurrent.futures.ThreadPoolExecutor( - max_workers=self.max_concurrent_jobs - ) as pool: - completed_jobs = list( - pool.map( - functools.partial(self._submit_and_poll_chunk, config=config), - chunks, - ) - ) - - logging.info('Batch annotate: all jobs completed. Parsing responses...') - - all_rows = [] - for batch_job, chunk_texts in zip(completed_jobs, chunks, strict=True): - chunk_rows = self._parse_job_responses(batch_job, schema) - - if len(chunk_rows) != len(chunk_texts): - raise ValueError( - f'Batch annotate: job {batch_job.name} got {len(chunk_rows)}' - f' results for {len(chunk_texts)} inputs.' - ) - - all_rows.extend(chunk_rows) - - return pd.DataFrame(all_rows) - - def generate(self, prompts: Sequence[str]) -> list[str]: - """Generate free-form text via google.genai. - - Args: - prompts: Fully constructed prompts. - - Returns: - List of exactly ``len(prompts)`` strings. Empty string on failure. - """ - client = self.client - results: list[str] = [] - for i, prompt in enumerate(prompts): - try: - response = client.models.generate_content( - model=self.model, - contents=prompt, - ) - results.append(response.text or '') - except Exception as e: # pylint: disable=broad-except - logging.warning( - 'Generation failed for prompt %d. Error: %s', i, e, exc_info=True - ) - results.append('') - return results - - -def _strip_markdown_fences(text): - """Strips markdown code fences from LLM output if present.""" - regex = r'^\s*```(?:json)?\s*\n(.*?)\n\s*```\s*$' - m = re.compile(regex, re.DOTALL).match(text) - return m.group(1).strip() if m else text.strip() - - -T = TypeVar('T') - - -def _call_with_retry( - func: Callable[[], T], - op_name: str, - max_retries: int = 10, - initial_delay: float = 5.0, -) -> T: # pyrefly: ignore[bad-return] - """Calls `func` with exponential backoff on exceptions.""" - delay = initial_delay - for attempt in range(1, max_retries + 1): - try: - return func() - except Exception as e: # pylint: disable=broad-except - if attempt == max_retries: - logging.error( - 'Batch %s failed after %d attempts.', op_name, max_retries - ) - raise - - sleep_time = delay + random.uniform(0, 5) - logging.warning( - 'Batch %s failed (attempt %d/%d): %s. Retrying in %.1f sec...', - op_name, - attempt, - max_retries, - e, - sleep_time, - ) - time.sleep(sleep_time) - delay *= 2 diff --git a/dpsynth/transformations.py b/dpsynth/transformations.py index 84312e2..48f99de 100644 --- a/dpsynth/transformations.py +++ b/dpsynth/transformations.py @@ -194,7 +194,7 @@ def create_discretize_transformation( ) bin_edges = np.r_[ - attribute_domain.min_value, + attribute_domain.exclusive_min_value, bin_edges, attribute_domain.max_value, ] @@ -266,7 +266,7 @@ def create_uniform_discretize_transformation( interval in the CategoricalAttribute and vice versa. """ bin_edges = np.linspace( - attribute_domain.min_value, + attribute_domain.exclusive_min_value, attribute_domain.max_value, num_bins + 1, ) diff --git a/tests/domain_test.py b/tests/domain_test.py index 012e769..e9cac5b 100644 --- a/tests/domain_test.py +++ b/tests/domain_test.py @@ -44,12 +44,18 @@ def test_invalid_out_of_domain_index(self): possible_values=['a', 'b'], out_of_domain_index=2 ) - def test_mixed_types_rejected(self): - with self.assertRaises(ValueError): - domain.CategoricalAttribute(possible_values=['a', 1]) + def test_none_rejected(self): with self.assertRaises(ValueError): domain.CategoricalAttribute(possible_values=[None, 'a']) + def test_non_categorical_type_rejected(self): + with self.assertRaises(ValueError): + domain.CategoricalAttribute(possible_values=['a', b'bytes']) + + def test_mixed_native_types_rejected(self): + with self.assertRaises(ValueError): + domain.CategoricalAttribute(possible_values=['', 1, True]) + def test_invalid_range(self): with self.assertRaises(ValueError): domain.NumericalAttribute(min_value=10, max_value=0) diff --git a/tests/text/bulk_inference_test.py b/tests/text/bulk_inference_test.py deleted file mode 100644 index d06622b..0000000 --- a/tests/text/bulk_inference_test.py +++ /dev/null @@ -1,377 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from typing import Literal -from unittest import mock - -from absl.testing import absltest -from dpsynth.text import bulk_inference -from dpsynth.text import prompts -import pandas as pd -import pydantic - -Topic = Literal['Science', 'Technology', 'Other'] -Complexity = Literal['Low', 'Medium', 'High'] - - -class SimpleFeatures(pydantic.BaseModel): - """Test feature schema.""" - - topic: Topic = pydantic.Field(description='The main topic.') - complexity: Complexity = pydantic.Field(description='Complexity level.') - - -class ModelNameTest(absltest.TestCase): - - def test_gemini_values(self): - self.assertEqual( - bulk_inference.ModelName.GEMINI_2_5_FLASH_LITE, 'gemini-2.5-flash-lite' - ) - self.assertEqual( - bulk_inference.ModelName.GEMINI_3_5_FLASH, 'gemini-3.5-flash' - ) - - def test_gemma_values(self): - self.assertEqual(bulk_inference.ModelName.GEMMA_4_27B, 'gemma-4-26b-a4b-it') - self.assertEqual(bulk_inference.ModelName.GEMMA_4_31B, 'gemma-4-31b-it') - - -class StripMarkdownFencesTest(absltest.TestCase): - - def test_strips_json_fence(self): - text = '```json\n{"topic": "Science"}\n```' - self.assertEqual( - bulk_inference._strip_markdown_fences(text), '{"topic": "Science"}' - ) - - def test_strips_plain_fence(self): - text = '```\n{"topic": "Science"}\n```' - self.assertEqual( - bulk_inference._strip_markdown_fences(text), '{"topic": "Science"}' - ) - - def test_noop_on_plain_json(self): - text = '{"topic": "Science"}' - self.assertEqual(bulk_inference._strip_markdown_fences(text), text) - - def test_strips_whitespace(self): - text = ' \n{"topic": "Science"}\n ' - self.assertEqual( - bulk_inference._strip_markdown_fences(text), '{"topic": "Science"}' - ) - - -@mock.patch('google.genai.Client', autospec=True) -class GenAIBackendAnnotateTest(absltest.TestCase): - """Tests for GenAIBackend.annotate.""" - - @staticmethod - def _create_inlined_response( - text: str | None = None, error: str | None = None - ) -> mock.MagicMock: - """Creates a mock inlined batch response.""" - mock_response = mock.MagicMock() - mock_response.text = text - return mock.MagicMock(error=error, response=mock_response if text else None) - - @staticmethod - def _create_mock_job( - state, - inlined_responses: list[mock.MagicMock] | None = None, - done_side_effect: list[bool] | None = None, - inlined_responses_side_effect: list[list[mock.MagicMock]] | None = None, - error: str | None = None, - name: str = 'job', - ) -> mock.MagicMock: - """Creates a mock batch job.""" - job = mock.MagicMock() - job.name = name - if done_side_effect is not None: - type(job).done = mock.PropertyMock(side_effect=done_side_effect) - else: - type(job).done = mock.PropertyMock(return_value=True) - job.state = state - job.error = error - - mock_dest = mock.MagicMock() - if inlined_responses_side_effect is not None: - type(mock_dest).inlined_responses = mock.PropertyMock( - side_effect=inlined_responses_side_effect - ) - elif inlined_responses is not None: - mock_dest.inlined_responses = inlined_responses - job.dest = mock_dest - return job - - def test_annotate_success(self, mock_client_cls): - mock_client = mock_client_cls.return_value - inlined_resp = self._create_inlined_response( - '{"topic": "Science", "complexity": "High"}' - ) - mock_batch_job = self._create_mock_job( - state=bulk_inference.types.JobState.JOB_STATE_SUCCEEDED, - inlined_responses=[inlined_resp, inlined_resp], - ) - - mock_client.batches.create.return_value = mock_batch_job - mock_client.batches.get.return_value = mock_batch_job - - backend = bulk_inference.GenAIBackend( - api_key='fake', poll_interval_seconds=0 - ) - df = backend.annotate(['text1', 'text2'], SimpleFeatures, 'Sys.') - self.assertLen(df, 2) - self.assertEqual(df.iloc[0]['topic'], 'Science') - self.assertEqual(df.iloc[1]['topic'], 'Science') - - def test_annotate_raises_on_length_mismatch(self, mock_client_cls): - mock_client = mock_client_cls.return_value - inlined_resp = self._create_inlined_response( - '{"topic": "Science", "complexity": "High"}' - ) - mock_batch_job = self._create_mock_job( - state=bulk_inference.types.JobState.JOB_STATE_SUCCEEDED, - inlined_responses=[inlined_resp], - ) - - mock_client.batches.create.return_value = mock_batch_job - mock_client.batches.get.return_value = mock_batch_job - - backend = bulk_inference.GenAIBackend( - api_key='fake', poll_interval_seconds=0 - ) - with self.assertRaisesRegex(ValueError, 'got 1 results for 2 inputs'): - backend.annotate(['text1', 'text2'], SimpleFeatures, 'Sys.') - - def test_annotate_fills_none_on_item_error(self, mock_client_cls): - mock_client = mock_client_cls.return_value - inlined_error = self._create_inlined_response(error='Failed item') - inlined_success = self._create_inlined_response( - '{"topic": "Science", "complexity": "High"}' - ) - - mock_batch_job = self._create_mock_job( - state=bulk_inference.types.JobState.JOB_STATE_SUCCEEDED, - inlined_responses=[inlined_error, inlined_success], - ) - - mock_client.batches.create.return_value = mock_batch_job - mock_client.batches.get.return_value = mock_batch_job - - backend = bulk_inference.GenAIBackend( - api_key='fake', poll_interval_seconds=0 - ) - df = backend.annotate(['text1', 'text2'], SimpleFeatures, 'Sys.') - self.assertLen(df, 2) - self.assertTrue(pd.isna(df.iloc[0]['topic'])) - self.assertEqual(df.iloc[1]['topic'], 'Science') - - def test_annotate_raises_on_failed_job(self, mock_client_cls): - mock_client = mock_client_cls.return_value - mock_batch_job = self._create_mock_job( - state=bulk_inference.types.JobState.JOB_STATE_FAILED, - error='Something went wrong', - ) - - mock_client.batches.create.return_value = mock_batch_job - mock_client.batches.get.return_value = mock_batch_job - - backend = bulk_inference.GenAIBackend( - api_key='fake', poll_interval_seconds=0 - ) - with self.assertRaisesRegex( - RuntimeError, - 'Batch job .* ended with state.* Error: Something went wrong', - ): - backend.annotate(['text1'], SimpleFeatures, 'Sys.') - - def test_annotate_respects_class_chunk_size(self, mock_client_cls): - mock_client = mock_client_cls.return_value - inlined_resp = self._create_inlined_response( - '{"topic": "Science", "complexity": "High"}' - ) - mock_batch_job = self._create_mock_job( - state=bulk_inference.types.JobState.JOB_STATE_SUCCEEDED, - inlined_responses=[inlined_resp, inlined_resp], - ) - - mock_client.batches.create.return_value = mock_batch_job - mock_client.batches.get.return_value = mock_batch_job - - backend = bulk_inference.GenAIBackend( - api_key='fake', poll_interval_seconds=0, chunk_size=2 - ) - # 4 texts, chunk_size=2 -> 2 chunks/jobs - backend.annotate(['t1', 't2', 't3', 't4'], SimpleFeatures, 'Sys.') - self.assertEqual(mock_client.batches.create.call_count, 2) - - def test_annotate_respects_max_concurrent_jobs(self, mock_client_cls): - mock_client = mock_client_cls.return_value - create_resp = GenAIBackendAnnotateTest._create_inlined_response - create_job = GenAIBackendAnnotateTest._create_mock_job - - jobs = [] - - def create_side_effect(model, src): - del model # Unused. - inlined_resp = create_resp('{"topic": "Science", "complexity": "High"}') - job = create_job( - state=bulk_inference.types.JobState.JOB_STATE_SUCCEEDED, - inlined_responses=[inlined_resp] * len(src), - done_side_effect=[False, True], - name=f'job-{len(jobs)}', - ) - jobs.append(job) - return job - - mock_client.batches.create.side_effect = create_side_effect - - def get_side_effect(name): - idx = int(name.split('-')[1]) - return jobs[idx] - - mock_client.batches.get.side_effect = get_side_effect - - backend = bulk_inference.GenAIBackend( - api_key='fake', - poll_interval_seconds=0, - chunk_size=2, - max_concurrent_jobs=1, - ) - backend.annotate(['t1', 't2', 't3', 't4', 't5'], SimpleFeatures, 'Sys.') - self.assertEqual(mock_client.batches.create.call_count, 3) - - -class GenAIBackendGenerateTest(absltest.TestCase): - - @mock.patch('google.genai.Client') - def test_generate_index_aligned_on_success(self, mock_client_cls): - mock_client = mock_client_cls.return_value - mock_response = mock.MagicMock() - mock_response.text = 'Generated text.' - mock_client.models.generate_content.return_value = mock_response - - backend = bulk_inference.GenAIBackend() - - results = backend.generate(['prompt1', 'prompt2']) - self.assertLen(results, 2) - self.assertEqual(results[0], 'Generated text.') - - @mock.patch('google.genai.Client') - def test_generate_fills_empty_on_failure(self, mock_client_cls): - mock_client = mock_client_cls.return_value - good = mock.MagicMock() - good.text = 'OK' - mock_client.models.generate_content.side_effect = [ - good, - RuntimeError('fail'), - good, - ] - - backend = bulk_inference.GenAIBackend() - - results = backend.generate(['a', 'b', 'c']) - self.assertLen(results, 3) - self.assertEqual(results[0], 'OK') - self.assertEqual(results[1], '') - self.assertEqual(results[2], 'OK') - - -class AnnotateFeaturesPromptTest(absltest.TestCase): - - def test_includes_dataset_description(self): - prompt = prompts.annotate_features_prompt( - dataset_description='News articles.', - dataclass=SimpleFeatures, - text='Some text.', - ) - self.assertIn('News articles', prompt) - - def test_includes_feature_names_and_descriptions(self): - prompt = prompts.annotate_features_prompt( - dataset_description='Test.', - dataclass=SimpleFeatures, - text='Hello.', - ) - self.assertIn('topic', prompt) - self.assertIn('complexity', prompt) - self.assertIn('The main topic', prompt) - - def test_includes_possible_values(self): - prompt = prompts.annotate_features_prompt( - dataset_description='Test.', - dataclass=SimpleFeatures, - text='Hello.', - ) - self.assertIn('Science', prompt) - self.assertIn('Technology', prompt) - - def test_includes_text_to_annotate(self): - prompt = prompts.annotate_features_prompt( - dataset_description='Test.', - dataclass=SimpleFeatures, - text='Quantum computing advances.', - ) - self.assertIn('Quantum computing advances', prompt) - - def test_template_placeholder(self): - prompt = prompts.annotate_features_prompt( - dataset_description='Test.', - dataclass=SimpleFeatures, - text='{{text}}', - ) - self.assertIn('{{text}}', prompt) - - -class ConditionalGenerationPromptTest(absltest.TestCase): - - def test_includes_target_features(self): - prompt = prompts.conditional_generation_prompt( - dataset_description='Test.', - target_features='- topic: Science\n- complexity: High', - formatting_requirements='Be concise.', - ) - self.assertIn('topic: Science', prompt) - self.assertIn('complexity: High', prompt) - - def test_includes_formatting_requirements(self): - prompt = prompts.conditional_generation_prompt( - dataset_description='Test.', - target_features='- topic: Tech', - formatting_requirements='Use bullet points.', - ) - self.assertIn('Use bullet points', prompt) - - def test_with_exemplar(self): - prompt = prompts.conditional_generation_prompt( - dataset_description='Test.', - target_features='- topic: Science', - formatting_requirements='Formal.', - exemplar=({'topic': 'Tech'}, 'Example tech article.'), - ) - self.assertIn('topic: Tech', prompt) - self.assertIn('Example tech article', prompt) - self.assertIn('Completed Example', prompt) - - def test_without_exemplar(self): - prompt = prompts.conditional_generation_prompt( - dataset_description='Test.', - target_features='- topic: Science', - formatting_requirements='Formal.', - ) - self.assertNotIn('Completed Example', prompt) - - -if __name__ == '__main__': - absltest.main() diff --git a/tests/transformations_test.py b/tests/transformations_test.py index 6ee84d8..25a505b 100644 --- a/tests/transformations_test.py +++ b/tests/transformations_test.py @@ -303,9 +303,9 @@ def test_binary_discretize(self): categorical, transformation = ( transformations.create_uniform_discretize_transformation(attr, 2) ) - self.assertEqual(categorical.possible_values, ['[0.0, 0.5]', '(0.5, 1.0]']) - self.assertEqual(transformation(0), '[0.0, 0.5]') - self.assertEqual(transformation(1), '(0.5, 1.0]') + self.assertEqual(categorical.possible_values, ['[0.0, 0.0]', '(0.0, 1.0]']) + self.assertEqual(transformation(0), '[0.0, 0.0]') + self.assertEqual(transformation(1), '(0.0, 1.0]') def test_applies_transformations_to_dataframe_drop_extra_columns(self): values = ['A', 'B', 'C']