diff --git a/nemo_curator/stages/base.py b/nemo_curator/stages/base.py index e0ae445a1b..3c55f15c98 100644 --- a/nemo_curator/stages/base.py +++ b/nemo_curator/stages/base.py @@ -112,6 +112,15 @@ def __init_subclass__(cls, **kwargs): msg = f"{cls.__name__} must not override '_batch_size'" raise TypeError(msg) + for attr in ("name", "resources", "batch_size"): + if isinstance(cls.__dict__.get(attr), property): + msg = ( + f"{cls.__name__} must not define '{attr}' as a @property. " + f"Use a plain class attribute or dataclass field instead, " + f"so that ProcessingStage.with_() can override it." + ) + raise TypeError(msg) + def num_workers(self) -> int | None: """Number of workers required. If None, then executor will determine the number of workers.""" return None diff --git a/nemo_curator/stages/synthetic/nemotron_cc/base.py b/nemo_curator/stages/synthetic/nemotron_cc/base.py index 6d93ebb394..bc929d3254 100644 --- a/nemo_curator/stages/synthetic/nemotron_cc/base.py +++ b/nemo_curator/stages/synthetic/nemotron_cc/base.py @@ -41,10 +41,7 @@ class BaseSyntheticStage(ProcessingStage[DocumentBatch, DocumentBatch]): client: AsyncLLMClient | LLMClient = None model_name: str = None generation_config: GenerationConfig | None = None - - @property - def name(self) -> str: - return "NemotronCCBaseStage" + name: str = "NemotronCCBaseStage" def __post_init__(self) -> None: self.is_async_client = isinstance(self.client, AsyncLLMClient) diff --git a/nemo_curator/stages/synthetic/nemotron_cc/nemotron_cc.py b/nemo_curator/stages/synthetic/nemotron_cc/nemotron_cc.py index 8964c4c972..c9f17fb838 100644 --- a/nemo_curator/stages/synthetic/nemotron_cc/nemotron_cc.py +++ b/nemo_curator/stages/synthetic/nemotron_cc/nemotron_cc.py @@ -64,10 +64,7 @@ class DiverseQAPostProcessingStage(ProcessingStage[DocumentBatch, DocumentBatch] tokenizer: AutoTokenizer | None = None prefix: str = "Here are the questions and answers based on the provided text:" max_num_pairs: int = 10 - - @property - def name(self) -> str: - return "DiverseQAPostProcessing" + name: str = "DiverseQAPostProcessing" def process(self, batch: DocumentBatch) -> DocumentBatch: df = batch.to_pandas() @@ -149,10 +146,7 @@ class KnowledgeListPostProcessingStage(ProcessingStage[DocumentBatch, DocumentBa """ input_field: str = "knowledge_list" - - @property - def name(self) -> str: - return "KnowledgeListPostProcessing" + name: str = "KnowledgeListPostProcessing" def process(self, batch: DocumentBatch) -> DocumentBatch: df = batch.to_pandas() diff --git a/tests/stages/common/test_base.py b/tests/stages/common/test_base.py index e1ee36cd31..937bc3d2f5 100644 --- a/tests/stages/common/test_base.py +++ b/tests/stages/common/test_base.py @@ -340,6 +340,46 @@ def inputs(self) -> tuple[list[str], list[str]]: def outputs(self) -> tuple[list[str], list[str]]: return [], [] + def test_name_property_decorator(self): + """Test that ProcessingStage raises an error if a derived class defines 'name' as a @property.""" + with pytest.raises(TypeError, match="must not define 'name' as a @property"): + + class MockStagePropertyName(ProcessingStage[MockTask, MockTask]): + @property + def name(self) -> str: + return "PropertyName" + + def process(self, task: MockTask) -> MockTask: + return task + + def test_resources_property_decorator(self): + """Test that ProcessingStage raises an error if a derived class defines 'resources' as a @property.""" + with pytest.raises(TypeError, match="must not define 'resources' as a @property"): + + class MockStagePropertyResources(ProcessingStage[MockTask, MockTask]): + name = "MockStagePropertyResources" + + @property + def resources(self) -> Resources: + return Resources(cpus=1.0) + + def process(self, task: MockTask) -> MockTask: + return task + + def test_batch_size_property_decorator(self): + """Test that ProcessingStage raises an error if a derived class defines 'batch_size' as a @property.""" + with pytest.raises(TypeError, match="must not define 'batch_size' as a @property"): + + class MockStagePropertyBatchSize(ProcessingStage[MockTask, MockTask]): + name = "MockStagePropertyBatchSize" + + @property + def batch_size(self) -> int: + return 1 + + def process(self, task: MockTask) -> MockTask: + return task + def test_nested_class_inheritance(self): """Test that nested class inheritance raises an error if a derived class overrides the _name, _resources, or _batch_size property.""" with pytest.raises(TypeError, match="MockStageNestedOverriddenName must not override '_name'"):