diff --git a/CHANGELOG.md b/CHANGELOG.md index 1333ce2..5c81145 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,8 +1,16 @@ +## Release v0.9.0 - TBD + +### Added + +- Added possibility to train models on custom datasets, without generating synthetic data. +- Added Italian `SpamDetection` model. +- Added the `TextSummarization` model. + ## Release v0.8.0 - March 26, 2026 ### Added -- Added German, Spanish Guardrail model and German, Spanish Spam Detection model. +- Added German, Spanish `Guardrail` model and German, Spanish `SpamDetection` model. ### Changed diff --git a/README.md b/README.md index fcad1f7..400cdc8 100644 --- a/README.md +++ b/README.md @@ -28,7 +28,7 @@

- Small Language Model Inference, Fine-Tuning and Observability. No GPU, no labeled data needed. + Small Language Model Inference, Fine-Tuning and Observability. No GPU, no labeled data needed.

--- @@ -61,21 +61,18 @@ Artifex changes that: At this time, Artifex supports the following models: -| Task | Language(s) | Description | Default Model | How to use | +| Task | Available Languages | Description | Default Model | How to use | |------|----------|-------------|---------------|------------| -| Guardrail | English | Flags unsafe, harmful, or off-topic messages. | [tanaos/tanaos-guardrail-v2](https://huggingface.co/tanaos/tanaos-guardrail-v2) | [Examples](https://docs.tanaos.com/artifex/guardrail/code-examples/) -| Guardrail | Spanish | Flags unsafe, harmful, or off-topic messages. | [tanaos/tanaos-guardrail-spanish](https://huggingface.co/tanaos/tanaos-guardrail-spanish) | [Examples](https://colab.research.google.com/#fileId=https%3A//huggingface.co/tanaos/tanaos-guardrail-spanish.ipynb) | -| Guardrail | German | Flags unsafe, harmful, or off-topic messages. | [tanaos/tanaos-guardrail-german](https://huggingface.co/tanaos/tanaos-guardrail-german) | [Examples](https://colab.research.google.com/#fileId=https%3A//huggingface.co/tanaos/tanaos-guardrail-german.ipynb) | +| Guardrail | English, German, Spanish | Flags unsafe, harmful, or off-topic messages. | [tanaos/tanaos-guardrail-v2](https://huggingface.co/tanaos/tanaos-guardrail-v2) (English version, see the page for the other languages) | [Examples](https://docs.tanaos.com/artifex/guardrail/code-examples/) | Intent Classification | English | Classifies user messages into predefined intent categories. | [tanaos/tanaos-intent-classifier-v1](https://huggingface.co/tanaos/tanaos-intent-classifier-v1) | [Examples](https://docs.tanaos.com/artifex/intent-classifier/code-examples/) | Reranker | English | Ranks a list of items or search results based on relevance to a query. | [cross-encoder/mmarco-mMiniLMv2-L12-H384-v1](https://huggingface.co/cross-encoder/mmarco-mMiniLMv2-L12-H384-v1) | [Examples](https://docs.tanaos.com/artifex/reranker/code-examples/) | Sentiment Analysis | English | Determines the sentiment (positive, negative, neutral) of a given text. | [tanaos/tanaos-sentiment-analysis-v1](https://huggingface.co/tanaos/tanaos-sentiment-analysis-v1) | [Examples](https://docs.tanaos.com/artifex/sentiment-analysis/code-examples/) | Emotion Detection | English | Identifies the emotion expressed in a given text. | [tanaos/tanaos-emotion-detection-v1](https://huggingface.co/tanaos/tanaos-emotion-detection-v1) | [Examples](https://docs.tanaos.com/artifex/emotion-detection/code-examples/) | Named Entity Recognition | English | Detects and classifies named entities in text (e.g., persons, organizations, locations). | [tanaos/tanaos-NER-v1](https://huggingface.co/tanaos/tanaos-NER-v1) | [Examples](https://docs.tanaos.com/artifex/named-entity-recognition/code-examples/) | Text Anonymization | English | Removes personally identifiable information (PII) from text. | [tanaos/tanaos-text-anonymizer-v1](https://huggingface.co/tanaos/tanaos-text-anonymizer-v1) | [Examples](https://docs.tanaos.com/artifex/text-anonymization/code-examples/) -| Spam Detection | English | Identifies whether a message is spam or not. | [tanaos/tanaos-spam-detection-v1](https://huggingface.co/tanaos/tanaos-spam-detection-v1) | [Examples](https://docs.tanaos.com/artifex/spam-detection/code-examples/) -| Spam Detection | Spanish | Identifies whether a message is spam or not. | [tanaos/tanaos-spam-detection-spanish](https://huggingface.co/tanaos/tanaos-spam-detection-spanish) | [Examples](https://colab.research.google.com/#fileId=https%3A//huggingface.co/tanaos/tanaos-spam-detection-spanish.ipynb) | -| Spam Detection | German | Identifies whether a message is spam or not. | [tanaos/tanaos-spam-detection-german](https://huggingface.co/tanaos/tanaos-spam-detection-german) | [Examples](https://colab.research.google.com/#fileId=https%3A//huggingface.co/tanaos/tanaos-spam-detection-german.ipynb) | +| Spam Detection | English, German, Spanish, Italian | Identifies whether a message is spam or not. | [tanaos/tanaos-spam-detection-v1](https://huggingface.co/tanaos/tanaos-spam-detection-v1) (English version, see the page for the other languages) | [Examples](https://docs.tanaos.com/artifex/spam-detection/code-examples/) | Topic Classification | English | Classifies text into predefined topics. | [tanaos/tanaos-topic-classification-v1](https://huggingface.co/tanaos/tanaos-topic-classification-v1) | [Examples](https://docs.tanaos.com/artifex/topic-classification/code-examples/) +| Text Summarization | English | Rewrites text in a more concise way. | [tanaos/tanaos-text-summarization-v1](https://huggingface.co/tanaos/tanaos-text-summarization-v1) | [Examples](https://docs.tanaos.com/artifex/text-summarization/code-examples/) Looking for models in other languages? Our [Enterprise License](#license-paid--enterprise-solutions) includes models in any language. Reach out at [info@tanaos.com](mailto:info@tanaos.com) for more details. diff --git a/artifex/__init__.py b/artifex/__init__.py index bba36ef..1b9f5b1 100644 --- a/artifex/__init__.py +++ b/artifex/__init__.py @@ -14,6 +14,7 @@ SpamDetection, TopicClassification from .models.named_entity_recognition import NamedEntityRecognition, TextAnonymization from .models.reranker import Reranker + from .models.text_summarization import TextSummarization from .config import config console.print(f"[green]✔ Initializing Artifex[/green]") @@ -56,6 +57,7 @@ def __init__(self, api_key: Optional[str] = None): self._text_anonymization = None self._spam_detection = None self._topic_classification = None + self._text_summarization = None def text_classification(self) -> ClassificationModel: """ @@ -160,12 +162,12 @@ def text_anonymization(self) -> TextAnonymization: return self._text_anonymization def spam_detection( - self, language: Literal["english", "spanish", "german"] = "english" + self, language: Literal["english", "spanish", "german", "italian"] = "english" ) -> SpamDetection: """ Lazy loads the SpamDetection instance. Args: - language (Literal["english", "spanish", "german"]): The language of the text data. + language (Literal["english", "spanish", "german", "italian"]): The language of the text data. Returns: SpamDetection: An instance of the SpamDetection class. """ @@ -187,4 +189,16 @@ def topic_classification(self) -> TopicClassification: if self._topic_classification is None: with console.status("Loading Topic Classification model..."): self._topic_classification = TopicClassification(synthex=self._synthex_client) - return self._topic_classification \ No newline at end of file + return self._topic_classification + + def text_summarization(self) -> TextSummarization: + """ + Lazy loads the TextSummarization instance. + Returns: + TextSummarization: An instance of the TextSummarization class. + """ + + if self._text_summarization is None: + with console.status("Loading Text Summarization model..."): + self._text_summarization = TextSummarization(synthex=self._synthex_client) + return self._text_summarization \ No newline at end of file diff --git a/artifex/config.py b/artifex/config.py index 05ee0c1..7b6064f 100644 --- a/artifex/config.py +++ b/artifex/config.py @@ -70,11 +70,17 @@ def DEFAULT_SYNTHEX_DATASET_NAME(self) -> str: NER_HF_BASE_MODEL: str = "tanaos/tanaos-NER-v1" NER_TOKENIZER_MAX_LENGTH: int = 256 NER_TAGNAME_MAX_LENGTH: int = 25 + + # Text Summarization Model + TEXT_SUMMARIZATION_HF_BASE_MODEL: str = "tanaos/tanaos-text-summarization-v1" + TEXT_SUMMARIZATION_MAX_INPUT_LENGTH: int = 512 + TEXT_SUMMARIZATION_MAX_TARGET_LENGTH: int = 128 # Spam Detection Models SPAM_DETECTION_ENGLISH_HF_BASE_MODEL: str = "tanaos/tanaos-spam-detection-v1" SPAM_DETECTION_SPANISH_HF_BASE_MODEL: str = "tanaos/tanaos-spam-detection-spanish" SPAM_DETECTION_GERMAN_HF_BASE_MODEL: str = "tanaos/tanaos-spam-detection-german" + SPAM_DETECTION_ITALIAN_HF_BASE_MODEL: str = "tanaos/tanaos-spam-detection-italian" SPAM_DETECTION_EVAL_DATASET: str = "hf://datasets/Deysi/spam-detection-dataset/data/train-00000-of-00001-daf190ce720b3dbb.parquet" # Topic Classification Model diff --git a/artifex/core/_hf_patches.py b/artifex/core/_hf_patches.py index d4d93be..0591bc0 100644 --- a/artifex/core/_hf_patches.py +++ b/artifex/core/_hf_patches.py @@ -1,4 +1,5 @@ -from transformers import Trainer, TrainerState, TrainingArguments, TrainerCallback, TrainerControl +from transformers import Trainer, TrainerState, TrainingArguments, TrainerCallback, TrainerControl, \ + Seq2SeqTrainer from transformers.trainer_utils import TrainOutput from typing import Any, Dict from rich.progress import Progress, BarColumn, TextColumn, TimeElapsedColumn, TaskID @@ -35,6 +36,32 @@ def silent_print(*a: Any, **k: Any) -> None: return super().train(*args, **kwargs) finally: builtins.print = orig_print + + +class SilentSeq2SeqTrainer(Seq2SeqTrainer): + """ + A regular transformers.Seq2SeqTrainer which prevents the tedious final training summary + dictionary from being printed to the console. + """ + + def train(self, *args: Any, **kwargs: Any) -> TrainOutput: + import builtins + orig_print = builtins.print + + def silent_print(*a: Any, **k: Any) -> None: + if ( + len(a) == 1 + and isinstance(a[0], dict) + and "train_runtime" in a[0] + ): + return + return orig_print(*a, **k) + + builtins.print = silent_print + try: + return super().train(*args, **kwargs) + finally: + builtins.print = orig_print class RichProgressCallback(TrainerCallback): diff --git a/artifex/models/__init__.py b/artifex/models/__init__.py index 2e3ed20..15033ee 100644 --- a/artifex/models/__init__.py +++ b/artifex/models/__init__.py @@ -7,6 +7,8 @@ from .reranker import Reranker +from .text_summarization import TextSummarization + __all__ = [ "ClassificationModel", "EmotionDetection", @@ -18,5 +20,6 @@ "Reranker", "SpamDetection", "TopicClassification", - "Guardrail" + "Guardrail", + "TextSummarization" ] \ No newline at end of file diff --git a/artifex/models/base_model.py b/artifex/models/base_model.py index 060c863..4f50579 100644 --- a/artifex/models/base_model.py +++ b/artifex/models/base_model.py @@ -131,16 +131,18 @@ def _synthetic_to_training_dataset(self, synthetic_dataset_path: str) -> Dataset @abstractmethod def _perform_train_pipeline( - self, user_instructions: ParsedModelInstructions, output_path: str, + self, user_instructions: Optional[ParsedModelInstructions], output_path: str, num_samples: int = config.DEFAULT_SYNTHEX_DATAPOINT_NUM, num_epochs: int = 3, train_datapoint_examples: Optional[list[dict[str, Any]]] = None, - device: Optional[int] = None + device: Optional[int] = None, + train_dataset_path: Optional[str] = None ) -> TrainOutput: f""" Perform the actual model training using the provided user instructions and training configuration. Args: - user_instructions (ParsedModelInstructions): A ParsedModelInstructions object containing user - instruction strings to be used for generating the training dataset. + user_instructions (Optional[ParsedModelInstructions]): A ParsedModelInstructions object containing + user instruction strings to be used for generating the training dataset. It can be None if + train_dataset_path is provided. output_path (Optional[str]): The directory path where training outputs and checkpoints will be saved. num_samples (Optional[int]): The number of synthetic datapoints to generate for training. Defaults to {config.DEFAULT_SYNTHEX_DATAPOINT_NUM}. @@ -149,6 +151,8 @@ def _perform_train_pipeline( the synthetic data generation. device (Optional[int]): The device to perform training on. If None, it will use the GPU if available, otherwise it will use the CPU. + train_dataset_path (Optional[str]): Path to an existing training dataset CSV file. If provided, + synthetic data generation is skipped and this dataset is used instead. Returns: TrainOutput: The output object containing training results and metrics. """ @@ -160,6 +164,7 @@ def train( self, language: str = "english", output_path: Optional[str] = None, num_samples: int = config.DEFAULT_SYNTHEX_DATAPOINT_NUM, num_epochs: int = 3, device: Optional[int] = None, disable_logging: Optional[bool] = False, + train_dataset_path: Optional[str] = None, *args: Any, **kwargs: Any, ) -> TrainOutput: f""" @@ -177,6 +182,9 @@ def train( device (Optional[int]): The device to perform training on. If None, it will use the GPU if available, otherwise it will use the CPU. disable_logging (Optional[bool]): Whether to disable logging during training. Defaults to False. + train_dataset_path (Optional[str]): Path to an existing training dataset CSV file. If provided, + synthetic data generation is skipped and this dataset is used instead. In this case, + synthetic data generation arguments (e.g., domain, classes, etc.) are optional. Returns: TrainOutput: The result of the training process, including metrics and model artifacts. """ @@ -400,44 +408,55 @@ def tokenize(example: dict[str, Sequence[str]]) -> BatchEncoding: return dataset.map(tokenize, batched=True) def _build_tokenized_train_ds( - self, user_instructions: ParsedModelInstructions, output_path: str, + self, user_instructions: Optional[ParsedModelInstructions], output_path: str, num_samples: int = config.DEFAULT_SYNTHEX_DATAPOINT_NUM, - train_datapoint_examples: Optional[list[dict[str, Any]]] = None + train_datapoint_examples: Optional[list[dict[str, Any]]] = None, + train_dataset_path: Optional[str] = None ) -> DatasetDict: """ Build a training dataset by generating synthetic data based on user-provided instructions and system instructions, then tokenize it. Args: - user_instructions (ParsedModelInstructions): A list of instructions, provided by the user, for - generating synthetic data. + user_instructions (Optional[ParsedModelInstructions]): A list of instructions, provided by the user, for + generating synthetic data. It can be None if train_dataset_path is provided. output_path (Optional[str]): The path where the generated synthetic data will be saved. num_samples (int): The number of training data samples to generate. train_datapoint_examples (Optional[list[dict[str, Any]]]): Examples of training datapoints to guide the synthetic data generation. + train_dataset_path (Optional[str]): Path to an existing training dataset CSV file. If provided, + synthetic data generation is skipped and this dataset is used instead. Returns: DatasetDict: The tokenized dataset ready for training. """ - output_dataset_path = get_dataset_output_path(output_path) + if train_dataset_path is not None: + # Use the provided training dataset path directly, skipping synthetic data generation. + output_dataset_path = train_dataset_path + else: + output_dataset_path = get_dataset_output_path(output_path) - # Build the data generation instructions by combining user instructions and system instructions - # NOTE: the system instructions MUST be prepended to the user instructions, as they provide - # context for the data generation. - full_instructions = self._get_data_gen_instr(user_instructions) + # Build the data generation instructions by combining user instructions and system instructions + # NOTE: the system instructions MUST be prepended to the user instructions, as they provide + # context for the data generation. + if user_instructions is None: + raise ValidationError( + message="User instructions must be provided when train_dataset_path is not provided." + ) + full_instructions = self._get_data_gen_instr(user_instructions) - # Generate synthetic data. - job_id = self._generate_synthetic_data( - schema_definition=self._synthetic_data_schema, - requirements=full_instructions, - output_path=output_dataset_path, - num_samples=num_samples, - examples=train_datapoint_examples - ) + # Generate synthetic data. + job_id = self._generate_synthetic_data( + schema_definition=self._synthetic_data_schema, + requirements=full_instructions, + output_path=output_dataset_path, + num_samples=num_samples, + examples=train_datapoint_examples + ) - # Await the completion of the synthetic data generation job. - self._await_data_generation( - get_status_fn=self._synthex.jobs.status, job_id=job_id - ) + # Await the completion of the synthetic data generation job. + self._await_data_generation( + get_status_fn=self._synthex.jobs.status, job_id=job_id + ) with console.status("Creating training dataset..."): # Once the job is complete, clean up the synthetic dataset (which may contain errors or @@ -454,10 +473,11 @@ def _build_tokenized_train_ds( return tokenized_dataset def _train_pipeline( - self, user_instructions: ParsedModelInstructions, output_path: Optional[str] = None, + self, user_instructions: Optional[ParsedModelInstructions], output_path: Optional[str] = None, num_samples: int = config.DEFAULT_SYNTHEX_DATAPOINT_NUM, num_epochs: int = 3, train_datapoint_examples: Optional[list[dict[str, Any]]] = None, - device: Optional[int] = None + device: Optional[int] = None, + train_dataset_path: Optional[str] = None ) -> TrainOutput: f""" NOTE: This method contains training-related logic that is common across all models. As such, it must @@ -469,8 +489,8 @@ def _train_pipeline( 3. Call the concrete `_perform_train_pipeline` method to perform the actual model training. 4. Print a success message with the model output path. Args: - user_instructions (list[str]): A list of user instruction strings to be used for generating the training - dataset. + user_instructions (Optional[ParsedModelInstructions]): A list of user instruction strings to be used for + generating the training dataset. It can be None if train_dataset_path is provided. output_path (Optional[str]): The directory path where training outputs and checkpoints will be saved. num_samples (Optional[int]): The number of synthetic datapoints to generate for training. Defaults to {config.DEFAULT_SYNTHEX_DATAPOINT_NUM}. @@ -479,6 +499,8 @@ def _train_pipeline( the synthetic data generation. device (Optional[int]): The device to perform training on. If None, it will use the GPU if available, otherwise it will use the CPU. + train_dataset_path (Optional[str]): Path to an existing training dataset CSV file. If provided, + synthetic data generation is skipped and this dataset is used instead. Returns: TrainOutput: The output object containing training results and metrics. """ @@ -510,7 +532,8 @@ def _train_pipeline( num_samples=num_samples, num_epochs=num_epochs, train_datapoint_examples=train_datapoint_examples, - device=device + device=device, + train_dataset_path=train_dataset_path ) # Get model output path based on the sanitized output path and print a success message diff --git a/artifex/models/classification/binary_classification/spam_detection/__init__.py b/artifex/models/classification/binary_classification/spam_detection/__init__.py index fc2c1c2..1dca54e 100644 --- a/artifex/models/classification/binary_classification/spam_detection/__init__.py +++ b/artifex/models/classification/binary_classification/spam_detection/__init__.py @@ -5,7 +5,7 @@ from ...classification_model import ClassificationModel -from artifex.core import auto_validate_methods, ParsedModelInstructions, track_training_calls +from artifex.core import auto_validate_methods, ParsedModelInstructions, track_training_calls, ValidationError from artifex.config import config @@ -15,19 +15,25 @@ class SpamDetection(ClassificationModel): A binary classification model for detecting spam content in emails, messages or other text data. """ - def __init__(self, synthex: Synthex, language: Literal["english", "spanish", "german"] = "english"): + def __init__( + self, synthex: Synthex, + language: Literal[ + "english", "spanish", "german", "italian" + ] = "english" + ): """ Initializes the class with a Synthex instance. Args: synthex (Synthex): An instance of the Synthex class to generate the synthetic data used to train the model. - language (Literal["english", "spanish", "german"]): The language of the text data. + language (Literal["english", "spanish", "german", "italian"]): The language of the text data. """ language_to_model = { "english": config.SPAM_DETECTION_ENGLISH_HF_BASE_MODEL, "spanish": config.SPAM_DETECTION_SPANISH_HF_BASE_MODEL, "german": config.SPAM_DETECTION_GERMAN_HF_BASE_MODEL, + "italian": config.SPAM_DETECTION_ITALIAN_HF_BASE_MODEL, } base_model = language_to_model[language] super().__init__(synthex, base_model_name=base_model) @@ -84,9 +90,10 @@ def _parse_user_instructions( @track_training_calls def train( - self, spam_content: list[str], language: str = "english", output_path: Optional[str] = None, + self, spam_content: Optional[list[str]] = None, language: str = "english", output_path: Optional[str] = None, num_samples: int = config.DEFAULT_SYNTHEX_DATAPOINT_NUM, num_epochs: int = 3, - device: Optional[int] = None, disable_logging: Optional[bool] = False + device: Optional[int] = None, disable_logging: Optional[bool] = False, + train_dataset_path: Optional[str] = None ) -> TrainOutput: f""" Overrides `ClassificationModel.train` to remove the `domain` and `classes` arguments and @@ -106,14 +113,20 @@ def train( TrainOutput: The output of the training process. """ + if train_dataset_path is None and spam_content is None: + raise ValidationError( + message="The `spam_content` parameter is required when `train_dataset_path` is not provided." + ) + user_instructions = self._parse_user_instructions( user_instructions=spam_content, language=language - ) + ) if spam_content is not None else None output: TrainOutput = self._train_pipeline( user_instructions=user_instructions, output_path=output_path, num_samples=num_samples, - num_epochs=num_epochs, device=device + num_epochs=num_epochs, device=device, + train_dataset_path=train_dataset_path ) return output \ No newline at end of file diff --git a/artifex/models/classification/classification_model.py b/artifex/models/classification/classification_model.py index ba6bb56..b79b037 100644 --- a/artifex/models/classification/classification_model.py +++ b/artifex/models/classification/classification_model.py @@ -167,10 +167,11 @@ def _synthetic_to_training_dataset(self, synthetic_dataset_path: str) -> Dataset return dataset def _perform_train_pipeline( - self, user_instructions: ParsedModelInstructions, output_path: str, + self, user_instructions: Optional[ParsedModelInstructions], output_path: str, num_samples: int = config.DEFAULT_SYNTHEX_DATAPOINT_NUM, num_epochs: int = 3, train_datapoint_examples: Optional[list[dict[str, Any]]] = None, - device: Optional[int] = None + device: Optional[int] = None, + train_dataset_path: Optional[str] = None ) -> TrainOutput: f""" Trains the model using the provided user instructions and training configuration. @@ -184,13 +185,16 @@ def _perform_train_pipeline( train_datapoint_examples (Optional[list[dict[str, Any]]]): Examples of training datapoints to guide the synthetic data generation. device (Optional[int]): The device to perform training on. If None, it will use the GPU if available, otherwise it will use the CPU. + train_dataset_path (Optional[str]): Path to an existing training dataset CSV file. If provided, + synthetic data generation is skipped and this dataset is used instead. Returns: TrainOutput: The output object containing training results and metrics. """ tokenized_dataset = self._build_tokenized_train_ds( user_instructions=user_instructions, output_path=output_path, - num_samples=num_samples + num_samples=num_samples, + train_dataset_path=train_dataset_path ) use_pin_memory = torch.cuda.is_available() or torch.backends.mps.is_available() @@ -230,18 +234,21 @@ def _perform_train_pipeline( @track_training_calls def train( - self, domain: str, classes: dict[str, str], language: str = "english", - output_path: Optional[str] = None, + self, domain: Optional[str] = None, classes: Optional[dict[str, str]] = None, + language: str = "english", output_path: Optional[str] = None, num_samples: int = config.DEFAULT_SYNTHEX_DATAPOINT_NUM, num_epochs: int = 3, - device: Optional[int] = None, disable_logging: Optional[bool] = False + device: Optional[int] = None, disable_logging: Optional[bool] = False, + train_dataset_path: Optional[str] = None ) -> TrainOutput: f""" Train the classification model using synthetic data generated by Synthex. Args: - domain (str): A description of the domain or context for which the model is being trained. - classes (dict[str, str]): A dictionary mapping class names to their descriptions. The keys + domain (Optional[str]): A description of the domain or context for which the model is being trained. + It is required if train_dataset_path is not provided. + classes (Optional[dict[str, str]]): A dictionary mapping class names to their descriptions. The keys (class names) must be string with no spaces and a maximum length of - {config.CLASSIFICATION_CLASS_NAME_MAX_LENGTH} characters. + {config.CLASSIFICATION_CLASS_NAME_MAX_LENGTH} characters. It is required if + train_dataset_path is not provided. language (str): The language in which the synthetic data should be generated. Defaults to "english". output_path (Optional[str]): The path where the generated synthetic data will be saved. num_samples (int): The number of training data samples to generate. @@ -249,19 +256,39 @@ def train( device (Optional[int]): The device to perform training on. If None, it will use the GPU if available, otherwise it will use the CPU. disable_logging (Optional[bool]): Whether to disable logging during training. Defaults to False. + train_dataset_path (Optional[str]): Path to an existing training dataset CSV file. If provided, + synthetic data generation is skipped and this dataset is used instead. Returns: TrainOutput: The output object containing training results and metrics. """ + if train_dataset_path is None and (domain is None or classes is None): + raise ValidationError( + message="The `domain` and `classes` parameters are required when `train_dataset_path` is not provided." + ) + # Validate class names, raise a ValidationError if any class name is invalid validated_classes: dict[str, str] = {} - for class_name, description in classes.items(): + if classes: + for class_name, description in classes.items(): + try: + validated_class_name = ClassificationClassName(class_name) + validated_classes[validated_class_name] = description + except ValueError: + raise ValidationError( + message=f"`classes` keys must be non-empty strings with no spaces and a maximum length of {config.CLASSIFICATION_CLASS_NAME_MAX_LENGTH} characters.", + ) + elif train_dataset_path: + # Infer classes from the dataset try: - validated_class_name = ClassificationClassName(class_name) - validated_classes[validated_class_name] = description - except ValueError: + df = pd.read_csv(train_dataset_path) + # The labels are in the last column + unique_labels = df.iloc[:, -1].unique() + for label in unique_labels: + validated_classes[str(label)] = f"Inferred class: {label}" + except Exception as e: raise ValidationError( - message=f"`classes` keys must be non-empty strings with no spaces and a maximum length of {config.CLASSIFICATION_CLASS_NAME_MAX_LENGTH} characters.", + message=f"Could not infer classes from the dataset: {str(e)}. Please provide the `classes` parameter." ) # Populate the labels property with the validated class names @@ -285,14 +312,15 @@ def train( user_instructions = self._parse_user_instructions( ClassificationInstructions( classes=validated_classes, - domain=domain, + domain=domain or "Inferred domain", language=language ) - ) + ) if domain is not None and classes is not None else None output: TrainOutput = self._train_pipeline( user_instructions=user_instructions, output_path=output_path, num_samples=num_samples, - num_epochs=num_epochs, device=device + num_epochs=num_epochs, device=device, + train_dataset_path=train_dataset_path ) return output diff --git a/artifex/models/classification/multi_class_classification/sentiment_analysis/__init__.py b/artifex/models/classification/multi_class_classification/sentiment_analysis/__init__.py index 7fdf418..05d3cf6 100644 --- a/artifex/models/classification/multi_class_classification/sentiment_analysis/__init__.py +++ b/artifex/models/classification/multi_class_classification/sentiment_analysis/__init__.py @@ -4,7 +4,7 @@ from ...classification_model import ClassificationModel -from artifex.core import auto_validate_methods, track_training_calls +from artifex.core import auto_validate_methods, track_training_calls, ValidationError from artifex.config import config @@ -35,9 +35,10 @@ def __init__(self, synthex: Synthex): @track_training_calls def train( - self, domain: str, classes: Optional[dict[str, str]] = None, language: str = "english", + self, domain: Optional[str] = None, classes: Optional[dict[str, str]] = None, language: str = "english", output_path: Optional[str] = None, num_samples: int = config.DEFAULT_SYNTHEX_DATAPOINT_NUM, - num_epochs: int = 3, device: Optional[int] = None, disable_logging: Optional[bool] = False + num_epochs: int = 3, device: Optional[int] = None, disable_logging: Optional[bool] = False, + train_dataset_path: Optional[str] = None ) -> TrainOutput: f""" Overrides `ClassificationModel.train()` to make the `classes` parameter optional. @@ -57,6 +58,11 @@ def train( TrainOutput: The output of the training process. """ + if train_dataset_path is None and domain is None: + raise ValidationError( + message="The `domain` parameter is required when `train_dataset_path` is not provided." + ) + if classes is None: classes = { "very_negative": "Text that expresses a very negative sentiment or strong dissatisfaction.", @@ -68,5 +74,6 @@ def train( return super().train( domain=domain, classes=classes, language=language, output_path=output_path, - num_samples=num_samples, num_epochs=num_epochs, device=device + num_samples=num_samples, num_epochs=num_epochs, device=device, + train_dataset_path=train_dataset_path ) diff --git a/artifex/models/classification/multi_label_classification/guardrail.py b/artifex/models/classification/multi_label_classification/guardrail.py index a7c1225..bcc865b 100644 --- a/artifex/models/classification/multi_label_classification/guardrail.py +++ b/artifex/models/classification/multi_label_classification/guardrail.py @@ -99,10 +99,11 @@ def _parse_user_instructions( @track_training_calls def train( - self, unsafe_categories: dict[str, str], language: str = "english", + self, unsafe_categories: Optional[dict[str, str]] = None, language: str = "english", output_path: Optional[str] = None, num_samples: int = config.DEFAULT_SYNTHEX_DATAPOINT_NUM, num_epochs: int = 3, - device: Optional[int] = None, disable_logging: Optional[bool] = False + device: Optional[int] = None, disable_logging: Optional[bool] = False, + train_dataset_path: Optional[str] = None ) -> TrainOutput: f""" Train the Guardrail model to detect multiple unsafe content categories in LLM outputs. @@ -128,6 +129,11 @@ def train( TrainOutput: The output of the training process. """ + if train_dataset_path is None and unsafe_categories is None: + raise ValidationError( + message="The `unsafe_categories` parameter is required when `train_dataset_path` is not provided." + ) + # Call the parent train method with the domain parameter return super().train( domain="LLM output safety and content moderation", @@ -137,7 +143,8 @@ def train( num_samples=num_samples, num_epochs=num_epochs, device=device, - disable_logging=disable_logging + disable_logging=disable_logging, + train_dataset_path=train_dataset_path ) @track_inference_calls diff --git a/artifex/models/classification/multi_label_classification/multi_label_classification_model.py b/artifex/models/classification/multi_label_classification/multi_label_classification_model.py index 0ae2f12..75b3779 100644 --- a/artifex/models/classification/multi_label_classification/multi_label_classification_model.py +++ b/artifex/models/classification/multi_label_classification/multi_label_classification_model.py @@ -227,16 +227,17 @@ def _synthetic_to_training_dataset(self, synthetic_dataset_path: str) -> Dataset return dataset def _perform_train_pipeline( - self, user_instructions: ParsedModelInstructions, output_path: str, + self, user_instructions: Optional[ParsedModelInstructions], output_path: str, num_samples: int = config.DEFAULT_SYNTHEX_DATAPOINT_NUM, num_epochs: int = 3, train_datapoint_examples: Optional[list[dict[str, Any]]] = None, - device: Optional[int] = None + device: Optional[int] = None, + train_dataset_path: Optional[str] = None ) -> TrainOutput: f""" Trains the multi-label model using the provided user instructions and training configuration. Args: - user_instructions (ParsedModelInstructions): A list of user instruction strings to be used for - generating the training dataset. + user_instructions (Optional[ParsedModelInstructions]): A list of user instruction strings to be used for + generating the training dataset. It can be None if train_dataset_path is provided. output_path (Optional[str]): The directory path where training outputs and checkpoints will be saved. num_samples (Optional[int]): The number of synthetic datapoints to generate for training. Defaults to {config.DEFAULT_SYNTHEX_DATAPOINT_NUM}. @@ -244,13 +245,16 @@ def _perform_train_pipeline( train_datapoint_examples (Optional[list[dict[str, Any]]]): Examples of training datapoints to guide the synthetic data generation. device (Optional[int]): The device to perform training on. If None, it will use the GPU if available, otherwise it will use the CPU. + train_dataset_path (Optional[str]): Path to an existing training dataset CSV file. If provided, + synthetic data generation is skipped and this dataset is used instead. Returns: TrainOutput: The output object containing training results and metrics. """ tokenized_dataset = self._build_tokenized_train_ds( user_instructions=user_instructions, output_path=output_path, - num_samples=num_samples + num_samples=num_samples, + train_dataset_path=train_dataset_path ) use_pin_memory = torch.cuda.is_available() or torch.backends.mps.is_available() @@ -291,18 +295,21 @@ def _perform_train_pipeline( @track_training_calls def train( - self, domain: str, labels: dict[str, str], language: str = "english", - output_path: Optional[str] = None, + self, domain: Optional[str] = None, labels: Optional[dict[str, str]] = None, + language: str = "english", output_path: Optional[str] = None, num_samples: int = config.DEFAULT_SYNTHEX_DATAPOINT_NUM, num_epochs: int = 3, - device: Optional[int] = None, disable_logging: Optional[bool] = False + device: Optional[int] = None, disable_logging: Optional[bool] = False, + train_dataset_path: Optional[str] = None ) -> TrainOutput: f""" Train the multi-label classification model using synthetic data generated by Synthex. Args: - domain (str): A description of the domain or context for which the model is being trained. - labels (dict[str, str]): A dictionary mapping label names to their descriptions. The keys + domain (Optional[str]): A description of the domain or context for which the model is being trained. + It is required if train_dataset_path is not provided. + labels (Optional[dict[str, str]]): A dictionary mapping label names to their descriptions. The keys (label names) must be strings with no spaces and a maximum length of - {config.CLASSIFICATION_CLASS_NAME_MAX_LENGTH} characters. + {config.CLASSIFICATION_CLASS_NAME_MAX_LENGTH} characters. It is required if + train_dataset_path is not provided. language (str): The language in which the synthetic data should be generated. Defaults to "english". output_path (Optional[str]): The path where the generated synthetic data will be saved. num_samples (int): The number of training data samples to generate. @@ -310,19 +317,54 @@ def train( device (Optional[int]): The device to perform training on. If None, it will use the GPU if available, otherwise it will use the CPU. disable_logging (Optional[bool]): Whether to disable logging during training. Defaults to False. + train_dataset_path (Optional[str]): Path to an existing training dataset CSV file. If provided, + synthetic data generation is skipped and this dataset is used instead. Returns: TrainOutput: The output object containing training results and metrics. """ + if train_dataset_path is None and (domain is None or labels is None): + raise ValidationError( + message="The `domain` and `labels` parameters are required when `train_dataset_path` is not provided." + ) + # Validate label names, raise a ValidationError if any label name is invalid validated_labels: dict[str, str] = {} - for label_name, description in labels.items(): + if labels: + for label_name, description in labels.items(): + try: + validated_label_name = ClassificationClassName(label_name) + validated_labels[validated_label_name] = description + except ValueError: + raise ValidationError( + message=f"`labels` keys must be non-empty strings with no spaces and a maximum length of {config.CLASSIFICATION_CLASS_NAME_MAX_LENGTH} characters.", + ) + elif train_dataset_path: + # Infer labels from the dataset try: - validated_label_name = ClassificationClassName(label_name) - validated_labels[validated_label_name] = description - except ValueError: + df = pd.read_csv(train_dataset_path) + # The labels are in the last column + unique_labels_str = df.iloc[:, -1].unique() + import ast + unique_labels = set() + for label_str in unique_labels_str: + if isinstance(label_str, str): + try: + parsed = ast.literal_eval(label_str) + if isinstance(parsed, list): + unique_labels.update(parsed) + else: + unique_labels.add(parsed) + except (ValueError, SyntaxError): + unique_labels.add(label_str) + else: + unique_labels.add(label_str) + + for label in unique_labels: + validated_labels[str(label)] = f"Inferred label: {label}" + except Exception as e: raise ValidationError( - message=f"`labels` keys must be non-empty strings with no spaces and a maximum length of {config.CLASSIFICATION_CLASS_NAME_MAX_LENGTH} characters.", + message=f"Could not infer labels from the dataset: {str(e)}. Please provide the `labels` parameter." ) # Populate the label names property @@ -346,14 +388,15 @@ def train( user_instructions = self._parse_user_instructions( ClassificationInstructions( classes=validated_labels, - domain=domain, + domain=domain or "Inferred domain", language=language ) - ) + ) if domain is not None and labels is not None else None output: TrainOutput = self._train_pipeline( user_instructions=user_instructions, output_path=output_path, num_samples=num_samples, - num_epochs=num_epochs, device=device + num_epochs=num_epochs, device=device, + train_dataset_path=train_dataset_path ) return output diff --git a/artifex/models/named_entity_recognition/named_entity_recognition.py b/artifex/models/named_entity_recognition/named_entity_recognition.py index dd4f986..d33fe95 100644 --- a/artifex/models/named_entity_recognition/named_entity_recognition.py +++ b/artifex/models/named_entity_recognition/named_entity_recognition.py @@ -288,15 +288,16 @@ def tokenize(example): return dataset.map(tokenize, batched=False) def _perform_train_pipeline( - self, user_instructions: ParsedModelInstructions, output_path: str, + self, user_instructions: Optional[ParsedModelInstructions], output_path: str, num_samples: int = config.DEFAULT_SYNTHEX_DATAPOINT_NUM, num_epochs: int = 3, train_datapoint_examples: Optional[list[dict[str, Any]]] = None, - device: Optional[int] = None + device: Optional[int] = None, + train_dataset_path: Optional[str] = None ) -> TrainOutput: f""" Trains the model using the provided user instructions and training configuration. Args: - user_instructions (ParsedModelInstructions): A list of user instruction strings to be used for generating the training dataset. + user_instructions (Optional[ParsedModelInstructions]): A list of user instruction strings to be used for generating the training dataset. It can be None if train_dataset_path is provided. output_path (Optional[str]): The directory path where training outputs and checkpoints will be saved. num_samples (Optional[int]): The number of synthetic datapoints to generate for training. Defaults to {config.DEFAULT_SYNTHEX_DATAPOINT_NUM}. @@ -305,13 +306,16 @@ def _perform_train_pipeline( datapoints to guide the synthetic data generation. device (Optional[int]): The device to perform training on. If None, it will use the GPU if available, otherwise it will use the CPU. + train_dataset_path (Optional[str]): Path to an existing training dataset CSV file. If provided, + synthetic data generation is skipped and this dataset is used instead. Returns: TrainOutput: The output object containing training results and metrics. """ tokenized_dataset = self._build_tokenized_train_ds( user_instructions=user_instructions, output_path=output_path, - num_samples=num_samples, train_datapoint_examples=train_datapoint_examples + num_samples=num_samples, train_datapoint_examples=train_datapoint_examples, + train_dataset_path=train_dataset_path ) use_pin_memory = torch.cuda.is_available() or torch.backends.mps.is_available() @@ -351,18 +355,21 @@ def _perform_train_pipeline( @track_training_calls def train( - self, named_entities: dict[str, str], domain: str, language: str = "english", - output_path: Optional[str] = None, num_samples: int = config.DEFAULT_SYNTHEX_DATAPOINT_NUM, + self, named_entities: Optional[dict[str, str]] = None, domain: Optional[str] = None, + language: str = "english", output_path: Optional[str] = None, + num_samples: int = config.DEFAULT_SYNTHEX_DATAPOINT_NUM, num_epochs: int = 3, train_datapoint_examples: Optional[list[dict[str, Any]]] = None, - device: Optional[int] = None, disable_logging: Optional[bool] = False + device: Optional[int] = None, disable_logging: Optional[bool] = False, + train_dataset_path: Optional[str] = None ) -> TrainOutput: f""" Train the classification model using synthetic data generated by Synthex. Args: - named_entities (dict[str, str]): A dictionary where keys are named entity tag names - and values are their descriptions. + named_entities (Optional[dict[str, str]]): A dictionary where keys are named entity tag names + and values are their descriptions. It is required if train_dataset_path is not provided. language (str): The language to use for generating the training dataset. Defaults to "english". - domain (str): The domain that the model will be specialized in. + domain (Optional[str]): The domain that the model will be specialized in. + It is required if train_dataset_path is not provided. output_path (Optional[str]): The path where the generated synthetic data will be saved. num_samples (int): The number of training data samples to generate. num_epochs (int): The number of epochs to train the model for. @@ -371,19 +378,53 @@ def train( device (Optional[int]): The device to perform training on. If None, it will use the GPU if available, otherwise it will use the CPU. disable_logging (Optional[bool]): Whether to disable logging during training. Defaults to False. + train_dataset_path (Optional[str]): Path to an existing training dataset CSV file. If provided, + synthetic data generation is skipped and this dataset is used instead. Returns: TrainOutput: The output object containing training results and metrics. """ + if train_dataset_path is None and (domain is None or named_entities is None): + raise ValidationError( + message="The `domain` and `named_entities` parameters are required when `train_dataset_path` is not provided." + ) + # Validate NER entity names, raise a ValidationError if any name is invalid validated_ner_instr: dict[str, str] = {} - for ner_name, description in named_entities.items(): + if named_entities: + for ner_name, description in named_entities.items(): + try: + validated_ner_name = NERTagName(ner_name) + validated_ner_instr[validated_ner_name] = description + except ValueError: + raise ValidationError( + message=f"`named_entities` keys must be non-empty strings with no spaces and a maximum length of {config.NER_TAGNAME_MAX_LENGTH} characters.", + ) + elif train_dataset_path: + # Infer entities from the dataset try: - validated_ner_name = NERTagName(ner_name) - validated_ner_instr[validated_ner_name] = description - except ValueError: + df = pd.read_csv(train_dataset_path) + # The labels are in the last column + unique_labels_str = df.iloc[:, -1].unique() + unique_tags = set() + for label_str in unique_labels_str: + if isinstance(label_str, str): + try: + parsed = ast.literal_eval(label_str) + if isinstance(parsed, list): + for tag in parsed: + if tag.startswith("B-") or tag.startswith("I-"): + unique_tags.add(tag[2:]) + elif tag != "O": + unique_tags.add(tag) + except (ValueError, SyntaxError): + pass + + for tag in unique_tags: + validated_ner_instr[tag] = f"Inferred entity: {tag}" + except Exception as e: raise ValidationError( - message=f"`named_entities` keys must be non-empty strings with no spaces and a maximum length of {config.NER_TAGNAME_MAX_LENGTH} characters.", + message=f"Could not infer entities from the dataset: {str(e)}. Please provide the `named_entities` parameter." ) # Populate the labels property with the validated class names; each class will @@ -411,14 +452,15 @@ def train( user_instructions = self._parse_user_instructions( NERInstructions( named_entity_tags=validated_ner_instr, - domain=domain, + domain=domain or "Inferred domain", language=language ) - ) + ) if domain is not None and named_entities is not None else None output: TrainOutput = self._train_pipeline( user_instructions=user_instructions, output_path=output_path, num_samples=num_samples, - num_epochs=num_epochs, train_datapoint_examples=train_datapoint_examples, device=device + num_epochs=num_epochs, train_datapoint_examples=train_datapoint_examples, device=device, + train_dataset_path=train_dataset_path ) return output diff --git a/artifex/models/named_entity_recognition/text_anonymization/__init__.py b/artifex/models/named_entity_recognition/text_anonymization/__init__.py index 472bad1..9b0f366 100644 --- a/artifex/models/named_entity_recognition/text_anonymization/__init__.py +++ b/artifex/models/named_entity_recognition/text_anonymization/__init__.py @@ -119,10 +119,11 @@ def __call__( @track_training_calls def train( - self, domain: str, pii_entities: dict[str, str], language: str = "english", + self, domain: Optional[str] = None, pii_entities: Optional[dict[str, str]] = None, language: str = "english", output_path: Optional[str] = None, num_samples: int = config.DEFAULT_SYNTHEX_DATAPOINT_NUM, num_epochs: int = 3, - device: Optional[int] = None, disable_logging: Optional[bool] = False + device: Optional[int] = None, disable_logging: Optional[bool] = False, + train_dataset_path: Optional[str] = None ) -> TrainOutput: """ Trains the Text Anonymization model. This method is identical to the @@ -144,18 +145,25 @@ def train( TrainOutput: The output of the training process. """ + if train_dataset_path is None and (domain is None or pii_entities is None): + raise ValidationError( + message="The `domain` and `pii_entities` parameters are required when `train_dataset_path` is not provided." + ) + # Validate PII entity names, raise a ValidationError if any name is invalid validated_ner_instr: dict[str, str] = {} - for ner_name, description in pii_entities.items(): - try: - validated_ner_name = NERTagName(ner_name) - validated_ner_instr[validated_ner_name] = description - except ValueError: - raise ValidationError( - message=f"`pii_entities` keys must be non-empty strings with no spaces and a maximum length of {config.NER_TAGNAME_MAX_LENGTH} characters.", - ) + if pii_entities: + for ner_name, description in pii_entities.items(): + try: + validated_ner_name = NERTagName(ner_name) + validated_ner_instr[validated_ner_name] = description + except ValueError: + raise ValidationError( + message=f"`pii_entities` keys must be non-empty strings with no spaces and a maximum length of {config.NER_TAGNAME_MAX_LENGTH} characters.", + ) return super().train( - named_entities=validated_ner_instr, domain=domain, language=language, output_path=output_path, - num_samples=num_samples, num_epochs=num_epochs, train_datapoint_examples=None, device=device + named_entities=validated_ner_instr or None, domain=domain, language=language, output_path=output_path, + num_samples=num_samples, num_epochs=num_epochs, train_datapoint_examples=None, device=device, + train_dataset_path=train_dataset_path ) \ No newline at end of file diff --git a/artifex/models/reranker/reranker.py b/artifex/models/reranker/reranker.py index 855cb0b..858b836 100644 --- a/artifex/models/reranker/reranker.py +++ b/artifex/models/reranker/reranker.py @@ -12,7 +12,7 @@ from ..base_model import BaseModel from artifex.core import auto_validate_methods, ParsedModelInstructions, track_inference_calls, \ - track_training_calls + track_training_calls, ValidationError from artifex.config import config from artifex.utils import get_model_output_path from artifex.core._hf_patches import SilentTrainer, RichProgressCallback @@ -176,15 +176,16 @@ def _synthetic_to_training_dataset(self, synthetic_dataset_path: str) -> Dataset return dataset def _perform_train_pipeline( - self, user_instructions: ParsedModelInstructions, output_path: str, + self, user_instructions: Optional[ParsedModelInstructions], output_path: str, num_samples: int = config.DEFAULT_SYNTHEX_DATAPOINT_NUM, num_epochs: int = 3, train_datapoint_examples: Optional[list[dict[str, Any]]] = None, - device: Optional[int] = None + device: Optional[int] = None, + train_dataset_path: Optional[str] = None ) -> TrainOutput: f""" Trains the model using the provided user instructions and training configuration. Args: - user_instructions (ParsedModelInstructions): A list of user instruction strings to be used for generating the training dataset. + user_instructions (Optional[ParsedModelInstructions]): A list of user instruction strings to be used for generating the training dataset. It can be None if train_dataset_path is provided. output_path (Optional[str]): The directory path where training outputs and checkpoints will be saved. num_samples (Optional[int]): The number of synthetic datapoints to generate for training. Defaults to {config.DEFAULT_SYNTHEX_DATAPOINT_NUM}. @@ -193,13 +194,16 @@ def _perform_train_pipeline( datapoints to guide the synthetic data generation. device (Optional[int]): The device to perform training on. If None, it will use the GPU if available, otherwise it will use the CPU. + train_dataset_path (Optional[str]): Path to an existing training dataset CSV file. If provided, + synthetic data generation is skipped and this dataset is used instead. Returns: TrainOutput: The output object containing training results and metrics. """ tokenized_dataset = self._build_tokenized_train_ds( user_instructions=user_instructions, output_path=output_path, - num_samples=num_samples, train_datapoint_examples=train_datapoint_examples + num_samples=num_samples, train_datapoint_examples=train_datapoint_examples, + train_dataset_path=train_dataset_path ) use_pin_memory = torch.cuda.is_available() or torch.backends.mps.is_available() @@ -239,15 +243,17 @@ def _perform_train_pipeline( @track_training_calls def train( - self, domain: str, language: str = "english", output_path: Optional[str] = None, + self, domain: Optional[str] = None, language: str = "english", output_path: Optional[str] = None, num_samples: int = config.DEFAULT_SYNTHEX_DATAPOINT_NUM, num_epochs: int = 3, train_datapoint_examples: Optional[list[dict[str, Any]]] = None, - device: Optional[int] = None, disable_logging: Optional[bool] = False + device: Optional[int] = None, disable_logging: Optional[bool] = False, + train_dataset_path: Optional[str] = None ) -> TrainOutput: f""" Train the classification model using synthetic data generated by Synthex. Args: - domain (str): The domain that the model will be specialized in. + domain (Optional[str]): The domain that the model will be specialized in. + It is required if train_dataset_path is not provided. language (str): The language of the training data. output_path (Optional[str]): The path where the generated synthetic data will be saved. num_samples (int): The number of training data samples to generate. @@ -257,20 +263,29 @@ def train( device (Optional[int]): The device to perform training on. If None, it will use the GPU if available, otherwise it will use the CPU. disable_logging (Optional[bool]): Whether to disable logging during training. Defaults to False. + train_dataset_path (Optional[str]): Path to an existing training dataset CSV file. If provided, + synthetic data generation is skipped and this dataset is used instead. Returns: TrainOutput: The output object containing training results and metrics. """ + if train_dataset_path is None and domain is None: + raise ValidationError( + message="The `domain` parameter is required when `train_dataset_path` is not provided." + ) + # Populate the domain property - self._domain = domain + if domain: + self._domain = domain # Turn domain into a list of strings, as expected by _train_pipeline - user_instructions = self._parse_user_instructions(domain, language) + user_instructions = self._parse_user_instructions(domain, language) if domain is not None else None output: TrainOutput = self._train_pipeline( user_instructions=user_instructions, output_path=output_path, num_samples=num_samples, num_epochs=num_epochs, train_datapoint_examples=train_datapoint_examples, - device=device + device=device, + train_dataset_path=train_dataset_path ) return output diff --git a/artifex/models/text_summarization/__init__.py b/artifex/models/text_summarization/__init__.py new file mode 100644 index 0000000..7c73565 --- /dev/null +++ b/artifex/models/text_summarization/__init__.py @@ -0,0 +1,6 @@ +from .text_summarization import TextSummarization + + +__all__ = [ + "TextSummarization" +] diff --git a/artifex/models/text_summarization/text_summarization.py b/artifex/models/text_summarization/text_summarization.py new file mode 100644 index 0000000..b33f91c --- /dev/null +++ b/artifex/models/text_summarization/text_summarization.py @@ -0,0 +1,362 @@ +from synthex import Synthex +from synthex.models import JobOutputSchemaDefinition +from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, PreTrainedTokenizerBase, \ + PreTrainedModel, Seq2SeqTrainingArguments, DataCollatorForSeq2Seq +from transformers.trainer_utils import TrainOutput +from typing import Optional, Union, Any, cast +import torch +import pandas as pd +from datasets import Dataset, DatasetDict +import os + +from ..base_model import BaseModel + +from artifex.core import auto_validate_methods, ParsedModelInstructions, track_inference_calls, \ + track_training_calls, ValidationError +from artifex.config import config +from artifex.utils import get_model_output_path +from artifex.core._hf_patches import SilentSeq2SeqTrainer, RichProgressCallback + + +@auto_validate_methods +class TextSummarization(BaseModel): + """ + A Text Summarization model takes a potentially long piece of text and returns a more + concise version of it, preserving only the key information. + """ + + def __init__(self, synthex: Synthex): + """ + Initializes the class with a Synthex instance. + Args: + synthex (Synthex): An instance of the Synthex class to generate the synthetic data used + to train the model. + """ + + super().__init__(synthex) + self._base_model_name_val: str = config.TEXT_SUMMARIZATION_HF_BASE_MODEL + self._synthetic_data_schema_val: JobOutputSchemaDefinition = { + "text": {"type": "string"}, + "summary": {"type": "string"}, + } + self._system_data_gen_instr_val: list[str] = [ + "The 'text' field should contain text belonging to the following domain: {domain}.", + "The 'text' field must be in the following language, and only this language: {language}.", + "The 'text' field should be at least 3 sentences long.", + "The 'summary' field should contain a concise summary of the 'text' field.", + "The 'summary' field must be in the following language, and only this language: {language}.", + "The 'summary' field should be significantly shorter than the 'text' field, capturing only the key information.", + "The 'summary' field should be 1 to 2 sentences long.", + ] + self._model_val: PreTrainedModel = AutoModelForSeq2SeqLM.from_pretrained( + self._base_model_name + ) + self._tokenizer_val: PreTrainedTokenizerBase = AutoTokenizer.from_pretrained( + self._base_model_name + ) + self._token_keys_val: list[str] = ["text", "summary"] + + @property + def _base_model_name(self) -> str: + return self._base_model_name_val + + @property + def _synthetic_data_schema(self) -> JobOutputSchemaDefinition: + return self._synthetic_data_schema_val + + @property + def _system_data_gen_instr(self) -> list[str]: + return self._system_data_gen_instr_val + + @property + def _token_keys(self) -> list[str]: + return self._token_keys_val + + def _parse_user_instructions( + self, user_instructions: str, language: str + ) -> ParsedModelInstructions: + """ + Convert the domain string passed by the user into a ParsedModelInstructions object. + Args: + user_instructions (str): The domain the summarization model will be specialized for. + language (str): The language to use for generating the training dataset. + Returns: + ParsedModelInstructions: The parsed instructions ready for data generation. + """ + + return ParsedModelInstructions( + user_instructions=[user_instructions], + language=language + ) + + def _get_data_gen_instr(self, user_instr: ParsedModelInstructions) -> list[str]: + """ + Generate data generation instructions by combining system instructions with user-provided + instructions. + Args: + user_instr (ParsedModelInstructions): A ParsedModelInstructions object containing user + instructions. + Returns: + list[str]: A list of fully formatted instructions for synthetic data generation. + """ + + return [ + instr.format( + domain=user_instr.user_instructions, + language=user_instr.language + ) for instr in self._system_data_gen_instr + ] + + def _post_process_synthetic_dataset(self, synthetic_dataset_path: str) -> None: + """ + Clean up the synthetic dataset for summarization training. Removes rows where: + - The 'text' field is missing, empty, or shorter than 50 characters. + - The 'summary' field is missing, empty, or shorter than 10 characters. + - The 'summary' field is not shorter than the 'text' field. + Args: + synthetic_dataset_path (str): The path to the synthetic dataset CSV file. + """ + + df = pd.read_csv(synthetic_dataset_path) + + df = df[df["text"].notna()] + df = df[df["text"].astype(str).str.strip() != ""] + df = df[df["text"].astype(str).str.strip().str.len() >= 50] + + df = df[df["summary"].notna()] + df = df[df["summary"].astype(str).str.strip() != ""] + df = df[df["summary"].astype(str).str.strip().str.len() >= 10] + + # Remove rows where the summary is not shorter than the text + df = df[df["summary"].astype(str).str.len() < df["text"].astype(str).str.len()] + + df.to_csv(synthetic_dataset_path, index=False) + + def _synthetic_to_training_dataset(self, synthetic_dataset_path: str) -> DatasetDict: + """ + Load the generated synthetic dataset from the specified path into a `datasets.Dataset` and + prepare it for training. + Args: + synthetic_dataset_path (str): The path to the synthetic dataset file. + Returns: + DatasetDict: A `datasets.DatasetDict` object containing the synthetic data, split into + training and validation sets. + """ + + dataset = cast(Dataset, Dataset.from_csv(synthetic_dataset_path)) + dataset = dataset.train_test_split(test_size=0.1) + + return dataset + + def _tokenize_dataset(self, dataset: DatasetDict, token_keys: list[str]) -> DatasetDict: + """ + Tokenize the dataset for seq2seq training: 'text' is tokenized as the encoder input and + 'summary' is tokenized as the decoder target (labels). Padding tokens in labels are + replaced with -100 so they are ignored during loss computation. + Args: + dataset (DatasetDict): The dataset to tokenize. + token_keys (list[str]): Unused; present for interface compatibility. + Returns: + DatasetDict: The tokenized dataset. + """ + + tokenizer = cast(Any, self._tokenizer) + + def tokenize(examples: dict[str, Any]) -> dict[str, Any]: + model_inputs = tokenizer( + examples["text"], + max_length=config.TEXT_SUMMARIZATION_MAX_INPUT_LENGTH, + truncation=True, + padding="max_length", + ) + labels = tokenizer( + text_target=examples["summary"], + max_length=config.TEXT_SUMMARIZATION_MAX_TARGET_LENGTH, + truncation=True, + padding="max_length", + ) + pad_id = tokenizer.pad_token_id + model_inputs["labels"] = [ + [(token_id if token_id != pad_id else -100) for token_id in ids] + for ids in labels["input_ids"] + ] + return model_inputs + + return dataset.map(tokenize, batched=True) + + def _perform_train_pipeline( + self, user_instructions: Optional[ParsedModelInstructions], output_path: str, + num_samples: int = config.DEFAULT_SYNTHEX_DATAPOINT_NUM, + num_epochs: int = 3, train_datapoint_examples: Optional[list[dict[str, Any]]] = None, + device: Optional[int] = None, + train_dataset_path: Optional[str] = None + ) -> TrainOutput: + f""" + Trains the model using the provided user instructions and training configuration. + Args: + user_instructions (Optional[ParsedModelInstructions]): A list of user instruction strings to be used for + generating the training dataset. It can be None if train_dataset_path is provided. + output_path (str): The directory path where training outputs and checkpoints will be saved. + num_samples (int): The number of synthetic datapoints to generate for training. Defaults to + {config.DEFAULT_SYNTHEX_DATAPOINT_NUM}. + num_epochs (int): The number of training epochs. Defaults to 3. + train_datapoint_examples (Optional[list[dict[str, Any]]]): Examples of training datapoints to + guide the synthetic data generation. + device (Optional[int]): The device to perform training on. If None, it will use the GPU + if available, otherwise it will use the CPU. + train_dataset_path (Optional[str]): Path to an existing training dataset CSV file. If provided, + synthetic data generation is skipped and this dataset is used instead. + Returns: + TrainOutput: The output object containing training results and metrics. + """ + + tokenized_dataset = self._build_tokenized_train_ds( + user_instructions=user_instructions, output_path=output_path, + num_samples=num_samples, train_datapoint_examples=train_datapoint_examples, + train_dataset_path=train_dataset_path + ) + + use_pin_memory = torch.cuda.is_available() or torch.backends.mps.is_available() + output_model_path = get_model_output_path(output_path) + + data_collator = DataCollatorForSeq2Seq( + tokenizer=cast(Any, self._tokenizer), + model=self._model, + label_pad_token_id=-100, + ) + + training_args = Seq2SeqTrainingArguments( + output_dir=output_model_path, + num_train_epochs=num_epochs, + per_device_train_batch_size=8, + per_device_eval_batch_size=8, + save_strategy="no", + logging_strategy="no", + report_to=[], + dataloader_pin_memory=use_pin_memory, + disable_tqdm=True, + use_cpu=self._should_disable_cuda(device), + predict_with_generate=True, + ) + + trainer = SilentSeq2SeqTrainer( + model=self._model, + args=training_args, + train_dataset=cast(Any, tokenized_dataset["train"]), + eval_dataset=cast(Any, tokenized_dataset["test"]), + data_collator=cast(Any, data_collator), + callbacks=[RichProgressCallback()], + ) + + train_output: TrainOutput = trainer.train() + trainer.save_model() + + training_args_path = os.path.join(output_model_path, "training_args.bin") + if os.path.exists(training_args_path): + os.remove(training_args_path) + + return train_output + + @track_training_calls + def train( + self, domain: Optional[str] = None, language: str = "english", output_path: Optional[str] = None, + num_samples: int = config.DEFAULT_SYNTHEX_DATAPOINT_NUM, num_epochs: int = 3, + train_datapoint_examples: Optional[list[dict[str, Any]]] = None, + device: Optional[int] = None, disable_logging: Optional[bool] = False, + train_dataset_path: Optional[str] = None + ) -> TrainOutput: + f""" + Train the text summarization model using synthetic data generated by Synthex. + Args: + domain (Optional[str]): The domain or topic that the summarization model will be specialized for. + It is required if train_dataset_path is not provided. + language (str): The language of the training data. Defaults to "english". + output_path (Optional[str]): The path where the trained model will be saved. + num_samples (int): The number of training data samples to generate. Defaults to + {config.DEFAULT_SYNTHEX_DATAPOINT_NUM}. + num_epochs (int): The number of epochs to train the model for. Defaults to 3. + train_datapoint_examples (Optional[list[dict[str, Any]]]): Examples of training datapoints + to guide the synthetic data generation. + device (Optional[int]): The device to perform training on. If None, it will use the GPU + if available, otherwise it will use the CPU. + disable_logging (Optional[bool]): Whether to disable logging during training. Defaults to False. + train_dataset_path (Optional[str]): Path to an existing training dataset CSV file. If provided, + synthetic data generation is skipped and this dataset is used instead. + Returns: + TrainOutput: The output object containing training results and metrics. + """ + + if train_dataset_path is None and domain is None: + raise ValidationError( + message="The `domain` parameter is required when `train_dataset_path` is not provided." + ) + + user_instructions = self._parse_user_instructions(domain, language) if domain else None + + output: TrainOutput = self._train_pipeline( + user_instructions=user_instructions, output_path=output_path, num_samples=num_samples, + num_epochs=num_epochs, train_datapoint_examples=train_datapoint_examples, device=device, + train_dataset_path=train_dataset_path + ) + + return output + + @track_inference_calls + def __call__( + self, text: Union[str, list[str]], device: Optional[int] = None, + disable_logging: Optional[bool] = False + ) -> list[str]: + """ + Summarize the input text. + Args: + text (Union[str, list[str]]): The input text or list of texts to summarize. + device (Optional[int]): The device to perform inference on. If None, it will use the GPU + if available, otherwise it will use the CPU. + disable_logging (Optional[bool]): Whether to disable logging during inference. Defaults to False. + Returns: + list[str]: A list of summaries, one per input text. + """ + + if self._model is None: + raise ValueError("Model not trained or loaded. Please call train() or load() first.") + + if device is None: + device = self._determine_default_device() + + if isinstance(text, str): + text = [text] + + torch_device = torch.device(f"cuda:{device}" if device >= 0 else "cpu") + self._model.to(torch_device) + + tokenizer = cast(Any, self._tokenizer) + inputs = tokenizer( + text, + max_length=config.TEXT_SUMMARIZATION_MAX_INPUT_LENGTH, + truncation=True, + padding=True, + return_tensors="pt", + ) + inputs = {k: v.to(torch_device) for k, v in inputs.items()} + + model = cast(Any, self._model) + with torch.no_grad(): + summary_ids = model.generate( + inputs["input_ids"], + attention_mask=inputs["attention_mask"], + max_length=config.TEXT_SUMMARIZATION_MAX_TARGET_LENGTH, + min_length=10, + num_beams=4, + ) + + return tokenizer.batch_decode(summary_ids, skip_special_tokens=True) + + def _load_model(self, model_path: str) -> None: + """ + Load a TextSummarization model from the specified path. + Args: + model_path (str): The path to the saved model. + """ + + self._model = AutoModelForSeq2SeqLM.from_pretrained(model_path) + self._tokenizer_val = AutoTokenizer.from_pretrained(model_path) diff --git a/tests/integration/text_summarization/__init__.py b/tests/integration/text_summarization/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/integration/text_summarization/test_ts__call__intgr.py b/tests/integration/text_summarization/test_ts__call__intgr.py new file mode 100644 index 0000000..029e80a --- /dev/null +++ b/tests/integration/text_summarization/test_ts__call__intgr.py @@ -0,0 +1,69 @@ +import pytest + +from artifex import Artifex + + +@pytest.mark.integration +def test__call__single_input_success( + artifex: Artifex +): + """ + Test the `__call__` method of TextSummarization with a single string input. + Ensure that: + - The return type is list[str]. + - The returned list contains exactly one string. + - The summary is shorter than the input text. + Args: + artifex (Artifex): The Artifex instance to be used for testing. + """ + + input_text = ( + "Cities grow faster than the systems designed to support them. " + "As populations increase, transportation, housing, and energy infrastructure " + "are stretched beyond their original limits. When planning fails to keep pace, " + "the result is congestion, higher living costs, and environmental strain. " + "Thoughtful urban design can reduce these pressures by prioritizing public transit, " + "mixed-use neighborhoods, and sustainable energy solutions." + ) + + out = artifex.text_summarization()(text=input_text, disable_logging=True) + + assert isinstance(out, list) + assert len(out) == 1 + assert isinstance(out[0], str) + assert len(out[0]) < len(input_text) + + +@pytest.mark.integration +def test__call__multiple_inputs_success( + artifex: Artifex +): + """ + Test the `__call__` method of TextSummarization with a list of inputs. + Ensure that: + - The return type is list[str]. + - The returned list has the same length as the input list. + Args: + artifex (Artifex): The Artifex instance to be used for testing. + """ + + input_texts = [ + ( + "Climate change is driven primarily by human activities, especially the burning " + "of fossil fuels. The resulting increase in greenhouse gas concentrations traps " + "more heat in the atmosphere, leading to rising temperatures worldwide. " + "These changes threaten ecosystems, sea levels, and human communities globally." + ), + ( + "Artificial intelligence is transforming industries from healthcare to finance. " + "Machine learning models can now diagnose diseases, detect fraud, and generate " + "creative content. Despite its promise, AI also raises concerns about job " + "displacement and algorithmic bias that society must address carefully." + ), + ] + + out = artifex.text_summarization()(text=input_texts, disable_logging=True) + + assert isinstance(out, list) + assert len(out) == len(input_texts) + assert all(isinstance(s, str) for s in out) diff --git a/tests/integration/text_summarization/test_ts_train_intgr.py b/tests/integration/text_summarization/test_ts_train_intgr.py new file mode 100644 index 0000000..99e4f96 --- /dev/null +++ b/tests/integration/text_summarization/test_ts_train_intgr.py @@ -0,0 +1,25 @@ +import pytest + +from artifex import Artifex + + +@pytest.mark.integration +def test_train_success( + artifex: Artifex, + output_folder: str +): + """ + Test the `train` method of the `TextSummarization` class. + Args: + artifex (Artifex): The Artifex instance to be used for testing. + output_folder (str): A temporary output folder path. + """ + + artifex.text_summarization().train( + domain="technology news", + num_samples=40, + num_epochs=1, + output_path=output_folder, + language="english", + disable_logging=True + ) diff --git a/tests/unit/base_model/test_bm_build_tokenized_train_ds.py b/tests/unit/base_model/test_bm_build_tokenized_train_ds.py index acf7bd9..7e348d9 100644 --- a/tests/unit/base_model/test_bm_build_tokenized_train_ds.py +++ b/tests/unit/base_model/test_bm_build_tokenized_train_ds.py @@ -90,7 +90,8 @@ def _synthetic_to_training_dataset(self, synthetic_dataset_path: str) -> Dataset def _perform_train_pipeline( self, user_instructions: ParsedModelInstructions, output_path: str, num_samples: int = config.DEFAULT_SYNTHEX_DATAPOINT_NUM, - num_epochs: int = 3, train_datapoint_examples: Optional[list[dict[str, Any]]] = None + num_epochs: int = 3, train_datapoint_examples: Optional[list[dict[str, Any]]] = None, + train_dataset_path: Optional[str] = None ) -> TrainOutput: return TrainOutput(global_step=100, training_loss=0.5, metrics={}) @@ -961,4 +962,200 @@ def test_build_tokenized_train_ds_with_multiple_examples( call_kwargs = mock_generate.call_args[1] assert call_kwargs['examples'] == examples - assert len(call_kwargs['examples']) == 3 \ No newline at end of file + assert len(call_kwargs['examples']) == 3 + + +@pytest.mark.unit +def test_build_tokenized_train_ds_skips_synthetic_generation_when_train_dataset_path_provided( + concrete_base_model: BaseModel, + mock_get_dataset_output_path: MagicMock, + mocker: MockerFixture +) -> None: + """ + Test that _build_tokenized_train_ds skips synthetic data generation when train_dataset_path is provided. + + Args: + concrete_base_model (BaseModel): The concrete BaseModel instance. + mock_get_dataset_output_path (MagicMock): Mocked get_dataset_output_path function. + mocker (MockerFixture): The pytest-mock fixture for mocking. + """ + + mock_get_instr = mocker.patch.object(concrete_base_model, '_get_data_gen_instr', return_value=["instr"]) + mock_generate = mocker.patch.object(concrete_base_model, '_generate_synthetic_data', return_value="job-id") + mock_await = mocker.patch.object(concrete_base_model, '_await_data_generation') + mocker.patch.object(concrete_base_model, '_post_process_synthetic_dataset') + mocker.patch.object(concrete_base_model, '_synthetic_to_training_dataset', return_value=DatasetDict()) + mocker.patch.object(concrete_base_model, '_tokenize_dataset', return_value=DatasetDict()) + + user_instructions = ParsedModelInstructions( + user_instructions=["instruction"], + language="english", + domain="test" + ) + + concrete_base_model._build_tokenized_train_ds( + user_instructions=user_instructions, + output_path="/path", + train_dataset_path="/existing/dataset.csv" + ) + + mock_get_instr.assert_not_called() + mock_generate.assert_not_called() + mock_await.assert_not_called() + + +@pytest.mark.unit +def test_build_tokenized_train_ds_uses_train_dataset_path_for_post_processing( + concrete_base_model: BaseModel, + mock_get_dataset_output_path: MagicMock, + mocker: MockerFixture +) -> None: + """ + Test that _build_tokenized_train_ds calls _post_process_synthetic_dataset with train_dataset_path. + + Args: + concrete_base_model (BaseModel): The concrete BaseModel instance. + mock_get_dataset_output_path (MagicMock): Mocked get_dataset_output_path function. + mocker (MockerFixture): The pytest-mock fixture for mocking. + """ + + mocker.patch.object(concrete_base_model, '_get_data_gen_instr', return_value=["instr"]) + mocker.patch.object(concrete_base_model, '_generate_synthetic_data', return_value="job-id") + mocker.patch.object(concrete_base_model, '_await_data_generation') + mock_post_process = mocker.patch.object(concrete_base_model, '_post_process_synthetic_dataset') + mocker.patch.object(concrete_base_model, '_synthetic_to_training_dataset', return_value=DatasetDict()) + mocker.patch.object(concrete_base_model, '_tokenize_dataset', return_value=DatasetDict()) + + user_instructions = ParsedModelInstructions( + user_instructions=["instruction"], + language="english", + domain="test" + ) + + concrete_base_model._build_tokenized_train_ds( + user_instructions=user_instructions, + output_path="/path", + train_dataset_path="/my/custom/dataset.csv" + ) + + mock_post_process.assert_called_once_with("/my/custom/dataset.csv") + + +@pytest.mark.unit +def test_build_tokenized_train_ds_uses_train_dataset_path_for_dataset_conversion( + concrete_base_model: BaseModel, + mock_get_dataset_output_path: MagicMock, + mocker: MockerFixture +) -> None: + """ + Test that _build_tokenized_train_ds calls _synthetic_to_training_dataset with train_dataset_path. + + Args: + concrete_base_model (BaseModel): The concrete BaseModel instance. + mock_get_dataset_output_path (MagicMock): Mocked get_dataset_output_path function. + mocker (MockerFixture): The pytest-mock fixture for mocking. + """ + + mocker.patch.object(concrete_base_model, '_get_data_gen_instr', return_value=["instr"]) + mocker.patch.object(concrete_base_model, '_generate_synthetic_data', return_value="job-id") + mocker.patch.object(concrete_base_model, '_await_data_generation') + mocker.patch.object(concrete_base_model, '_post_process_synthetic_dataset') + mock_to_dataset = mocker.patch.object( + concrete_base_model, '_synthetic_to_training_dataset', return_value=DatasetDict() + ) + mocker.patch.object(concrete_base_model, '_tokenize_dataset', return_value=DatasetDict()) + + user_instructions = ParsedModelInstructions( + user_instructions=["instruction"], + language="english", + domain="test" + ) + + concrete_base_model._build_tokenized_train_ds( + user_instructions=user_instructions, + output_path="/path", + train_dataset_path="/my/custom/dataset.csv" + ) + + mock_to_dataset.assert_called_once_with("/my/custom/dataset.csv") + + +@pytest.mark.unit +def test_build_tokenized_train_ds_still_tokenizes_when_train_dataset_path_provided( + concrete_base_model: BaseModel, + mock_get_dataset_output_path: MagicMock, + mocker: MockerFixture +) -> None: + """ + Test that _build_tokenized_train_ds still calls _tokenize_dataset when train_dataset_path is provided. + + Args: + concrete_base_model (BaseModel): The concrete BaseModel instance. + mock_get_dataset_output_path (MagicMock): Mocked get_dataset_output_path function. + mocker (MockerFixture): The pytest-mock fixture for mocking. + """ + + mocker.patch.object(concrete_base_model, '_get_data_gen_instr', return_value=["instr"]) + mocker.patch.object(concrete_base_model, '_generate_synthetic_data', return_value="job-id") + mocker.patch.object(concrete_base_model, '_await_data_generation') + mocker.patch.object(concrete_base_model, '_post_process_synthetic_dataset') + + mock_dataset = DatasetDict({ + "train": Dataset.from_dict({"text": ["sample"]}), + "test": Dataset.from_dict({"text": ["test"]}) + }) + mocker.patch.object(concrete_base_model, '_synthetic_to_training_dataset', return_value=mock_dataset) + mock_tokenize = mocker.patch.object( + concrete_base_model, '_tokenize_dataset', return_value=DatasetDict() + ) + + user_instructions = ParsedModelInstructions( + user_instructions=["instruction"], + language="english", + domain="test" + ) + + concrete_base_model._build_tokenized_train_ds( + user_instructions=user_instructions, + output_path="/path", + train_dataset_path="/my/custom/dataset.csv" + ) + + mock_tokenize.assert_called_once_with(mock_dataset, concrete_base_model._token_keys) + + +@pytest.mark.unit +def test_build_tokenized_train_ds_skips_get_dataset_output_path_when_train_dataset_path_provided( + concrete_base_model: BaseModel, + mock_get_dataset_output_path: MagicMock, + mocker: MockerFixture +) -> None: + """ + Test that _build_tokenized_train_ds does NOT call get_dataset_output_path when train_dataset_path is provided. + + Args: + concrete_base_model (BaseModel): The concrete BaseModel instance. + mock_get_dataset_output_path (MagicMock): Mocked get_dataset_output_path function. + mocker (MockerFixture): The pytest-mock fixture for mocking. + """ + + mocker.patch.object(concrete_base_model, '_get_data_gen_instr', return_value=["instr"]) + mocker.patch.object(concrete_base_model, '_generate_synthetic_data', return_value="job-id") + mocker.patch.object(concrete_base_model, '_await_data_generation') + mocker.patch.object(concrete_base_model, '_post_process_synthetic_dataset') + mocker.patch.object(concrete_base_model, '_synthetic_to_training_dataset', return_value=DatasetDict()) + mocker.patch.object(concrete_base_model, '_tokenize_dataset', return_value=DatasetDict()) + + user_instructions = ParsedModelInstructions( + user_instructions=["instruction"], + language="english", + domain="test" + ) + + concrete_base_model._build_tokenized_train_ds( + user_instructions=user_instructions, + output_path="/path", + train_dataset_path="/existing/dataset.csv" + ) + + mock_get_dataset_output_path.assert_not_called() \ No newline at end of file diff --git a/tests/unit/base_model/test_bm_train_pipeline.py b/tests/unit/base_model/test_bm_train_pipeline.py index db5f207..57a24d6 100644 --- a/tests/unit/base_model/test_bm_train_pipeline.py +++ b/tests/unit/base_model/test_bm_train_pipeline.py @@ -81,7 +81,8 @@ def _perform_train_pipeline( self, user_instructions: ParsedModelInstructions, output_path: str, num_samples: int = config.DEFAULT_SYNTHEX_DATAPOINT_NUM, num_epochs: int = 3, train_datapoint_examples: Optional[list[dict[str, Any]]] = None, - device: Optional[int] = None + device: Optional[int] = None, + train_dataset_path: Optional[str] = None ) -> TrainOutput: return TrainOutput(global_step=100, training_loss=0.5, metrics={}) @@ -277,7 +278,8 @@ def test_train_pipeline_calls_perform_train_pipeline( num_samples=100, num_epochs=5, train_datapoint_examples=None, - device=None + device=None, + train_dataset_path=None ) @@ -743,4 +745,71 @@ def test_train_pipeline_with_empty_train_datapoint_examples( train_datapoint_examples=[] ) - assert isinstance(result, TrainOutput) \ No newline at end of file + assert isinstance(result, TrainOutput) + + +@pytest.mark.unit +def test_train_pipeline_passes_train_dataset_path_to_perform_train_pipeline( + concrete_base_model: BaseModel, + mocker: MockerFixture +) -> None: + """ + Test that _train_pipeline passes train_dataset_path to _perform_train_pipeline. + + Args: + concrete_base_model (BaseModel): The concrete BaseModel instance. + mocker (MockerFixture): The pytest-mock fixture for mocking. + """ + + mocker.patch('os.path.exists', return_value=False) + mocker.patch('artifex.models.base_model.get_model_output_path', return_value="/model/path") + + mock_perform = mocker.spy(concrete_base_model, '_perform_train_pipeline') + + user_instructions = ParsedModelInstructions( + user_instructions=["test"], + language="english", + domain="test" + ) + + concrete_base_model._train_pipeline( + user_instructions=user_instructions, + output_path="/path", + train_dataset_path="/existing/dataset.csv" + ) + + call_kwargs = mock_perform.call_args[1] + assert call_kwargs['train_dataset_path'] == "/existing/dataset.csv" + + +@pytest.mark.unit +def test_train_pipeline_passes_none_train_dataset_path_when_not_specified( + concrete_base_model: BaseModel, + mocker: MockerFixture +) -> None: + """ + Test that _train_pipeline passes train_dataset_path=None when not specified. + + Args: + concrete_base_model (BaseModel): The concrete BaseModel instance. + mocker (MockerFixture): The pytest-mock fixture for mocking. + """ + + mocker.patch('os.path.exists', return_value=False) + mocker.patch('artifex.models.base_model.get_model_output_path', return_value="/model/path") + + mock_perform = mocker.spy(concrete_base_model, '_perform_train_pipeline') + + user_instructions = ParsedModelInstructions( + user_instructions=["test"], + language="english", + domain="test" + ) + + concrete_base_model._train_pipeline( + user_instructions=user_instructions, + output_path="/path" + ) + + call_kwargs = mock_perform.call_args[1] + assert call_kwargs['train_dataset_path'] is None \ No newline at end of file diff --git a/tests/unit/classification_model/test_cm_perform_train_pipeline.py b/tests/unit/classification_model/test_cm_perform_train_pipeline.py index 063aa04..151842a 100644 --- a/tests/unit/classification_model/test_cm_perform_train_pipeline.py +++ b/tests/unit/classification_model/test_cm_perform_train_pipeline.py @@ -204,7 +204,8 @@ def test_perform_train_pipeline_calls_build_tokenized_train_ds( mock_build.assert_called_once_with( user_instructions=user_instructions, output_path="/output", - num_samples=200 + num_samples=200, + train_dataset_path=None ) diff --git a/tests/unit/multi_label_classification_model/test_mlcm_get_data_gen_instr.py b/tests/unit/multi_label_classification_model/test_mlcm_get_data_gen_instr.py index 4592b98..35707d0 100644 --- a/tests/unit/multi_label_classification_model/test_mlcm_get_data_gen_instr.py +++ b/tests/unit/multi_label_classification_model/test_mlcm_get_data_gen_instr.py @@ -39,7 +39,20 @@ def mock_tokenizer(mocker: MockerFixture) -> MagicMock: @pytest.fixture -def mlcm_instance(mock_synthex: MagicMock, mock_tokenizer: MagicMock) -> MultiLabelClassificationModel: +def mock_model(mocker: MockerFixture) -> MagicMock: + """ + Fixture that provides a mock model and patches AutoModelForSequenceClassification.from_pretrained. + """ + mock_mdl = MagicMock() + mocker.patch( + 'artifex.models.classification.multi_label_classification.multi_label_classification_model.AutoModelForSequenceClassification.from_pretrained', + return_value=mock_mdl + ) + return mock_mdl + + +@pytest.fixture +def mlcm_instance(mock_synthex: MagicMock, mock_tokenizer: MagicMock, mock_model: MagicMock) -> MultiLabelClassificationModel: """ Fixture that provides a MultiLabelClassificationModel instance with preset label names. diff --git a/tests/unit/multi_label_classification_model/test_mlcm_load_model.py b/tests/unit/multi_label_classification_model/test_mlcm_load_model.py index beb34a7..74c10e7 100644 --- a/tests/unit/multi_label_classification_model/test_mlcm_load_model.py +++ b/tests/unit/multi_label_classification_model/test_mlcm_load_model.py @@ -39,7 +39,20 @@ def mock_tokenizer(mocker: MockerFixture) -> MagicMock: @pytest.fixture -def mlcm_instance(mock_synthex: MagicMock, mock_tokenizer: MagicMock) -> MultiLabelClassificationModel: +def mock_model(mocker: MockerFixture) -> MagicMock: + """ + Fixture that provides a mock model and patches AutoModelForSequenceClassification.from_pretrained. + """ + mock_mdl = MagicMock() + mocker.patch( + 'artifex.models.classification.multi_label_classification.multi_label_classification_model.AutoModelForSequenceClassification.from_pretrained', + return_value=mock_mdl + ) + return mock_mdl + + +@pytest.fixture +def mlcm_instance(mock_synthex: MagicMock, mock_tokenizer: MagicMock, mock_model: MagicMock) -> MultiLabelClassificationModel: """ Fixture that provides a MultiLabelClassificationModel instance. diff --git a/tests/unit/multi_label_classification_model/test_mlcm_parse_user_instructions.py b/tests/unit/multi_label_classification_model/test_mlcm_parse_user_instructions.py index 122e21c..cd763cf 100644 --- a/tests/unit/multi_label_classification_model/test_mlcm_parse_user_instructions.py +++ b/tests/unit/multi_label_classification_model/test_mlcm_parse_user_instructions.py @@ -39,7 +39,20 @@ def mock_tokenizer(mocker: MockerFixture) -> MagicMock: @pytest.fixture -def mlcm_instance(mock_synthex: MagicMock, mock_tokenizer: MagicMock) -> MultiLabelClassificationModel: +def mock_model(mocker: MockerFixture) -> MagicMock: + """ + Fixture that provides a mock model and patches AutoModelForSequenceClassification.from_pretrained. + """ + mock_mdl = MagicMock() + mocker.patch( + 'artifex.models.classification.multi_label_classification.multi_label_classification_model.AutoModelForSequenceClassification.from_pretrained', + return_value=mock_mdl + ) + return mock_mdl + + +@pytest.fixture +def mlcm_instance(mock_synthex: MagicMock, mock_tokenizer: MagicMock, mock_model: MagicMock) -> MultiLabelClassificationModel: """ Fixture that provides a MultiLabelClassificationModel instance. diff --git a/tests/unit/multi_label_classification_model/test_mlcm_post_process_synthetic_dataset.py b/tests/unit/multi_label_classification_model/test_mlcm_post_process_synthetic_dataset.py index b6e1a36..2cd7eda 100644 --- a/tests/unit/multi_label_classification_model/test_mlcm_post_process_synthetic_dataset.py +++ b/tests/unit/multi_label_classification_model/test_mlcm_post_process_synthetic_dataset.py @@ -39,7 +39,20 @@ def mock_tokenizer(mocker: MockerFixture) -> MagicMock: @pytest.fixture -def mlcm_instance(mock_synthex: MagicMock, mock_tokenizer: MagicMock) -> MultiLabelClassificationModel: +def mock_model(mocker: MockerFixture) -> MagicMock: + """ + Fixture that provides a mock model and patches AutoModelForSequenceClassification.from_pretrained. + """ + mock_mdl = MagicMock() + mocker.patch( + 'artifex.models.classification.multi_label_classification.multi_label_classification_model.AutoModelForSequenceClassification.from_pretrained', + return_value=mock_mdl + ) + return mock_mdl + + +@pytest.fixture +def mlcm_instance(mock_synthex: MagicMock, mock_tokenizer: MagicMock, mock_model: MagicMock) -> MultiLabelClassificationModel: """ Fixture that provides a MultiLabelClassificationModel instance with preset label names. diff --git a/tests/unit/multi_label_classification_model/test_mlcm_synthetic_to_training_dataset.py b/tests/unit/multi_label_classification_model/test_mlcm_synthetic_to_training_dataset.py index 9ded455..87a35f0 100644 --- a/tests/unit/multi_label_classification_model/test_mlcm_synthetic_to_training_dataset.py +++ b/tests/unit/multi_label_classification_model/test_mlcm_synthetic_to_training_dataset.py @@ -40,7 +40,20 @@ def mock_tokenizer(mocker: MockerFixture) -> MagicMock: @pytest.fixture -def mlcm_instance(mock_synthex: MagicMock, mock_tokenizer: MagicMock) -> MultiLabelClassificationModel: +def mock_model(mocker: MockerFixture) -> MagicMock: + """ + Fixture that provides a mock model and patches AutoModelForSequenceClassification.from_pretrained. + """ + mock_mdl = MagicMock() + mocker.patch( + 'artifex.models.classification.multi_label_classification.multi_label_classification_model.AutoModelForSequenceClassification.from_pretrained', + return_value=mock_mdl + ) + return mock_mdl + + +@pytest.fixture +def mlcm_instance(mock_synthex: MagicMock, mock_tokenizer: MagicMock, mock_model: MagicMock) -> MultiLabelClassificationModel: """ Fixture that provides a MultiLabelClassificationModel instance with preset label names. diff --git a/tests/unit/named_entity_recognition/test_ner_perform_train_pipeline.py b/tests/unit/named_entity_recognition/test_ner_perform_train_pipeline.py index f371f36..7ce4ca8 100644 --- a/tests/unit/named_entity_recognition/test_ner_perform_train_pipeline.py +++ b/tests/unit/named_entity_recognition/test_ner_perform_train_pipeline.py @@ -205,7 +205,8 @@ def test_perform_train_pipeline_calls_build_tokenized_train_ds( user_instructions=user_instructions, output_path="/output", num_samples=200, - train_datapoint_examples=None + train_datapoint_examples=None, + train_dataset_path=None ) @@ -253,7 +254,8 @@ def test_perform_train_pipeline_calls_build_tokenized_train_ds_with_examples( user_instructions=user_instructions, output_path="/output", num_samples=200, - train_datapoint_examples=examples + train_datapoint_examples=examples, + train_dataset_path=None ) diff --git a/tests/unit/reranker/test_rr_perform_train_pipeline.py b/tests/unit/reranker/test_rr_perform_train_pipeline.py index f9c63c1..039c22c 100644 --- a/tests/unit/reranker/test_rr_perform_train_pipeline.py +++ b/tests/unit/reranker/test_rr_perform_train_pipeline.py @@ -204,7 +204,8 @@ def test_perform_train_pipeline_calls_build_tokenized_train_ds( user_instructions=user_instructions, output_path="/output", num_samples=200, - train_datapoint_examples=None + train_datapoint_examples=None, + train_dataset_path=None ) diff --git a/tests/unit/reranker/test_rr_train.py b/tests/unit/reranker/test_rr_train.py index 1852de2..6ce5d70 100644 --- a/tests/unit/reranker/test_rr_train.py +++ b/tests/unit/reranker/test_rr_train.py @@ -166,7 +166,8 @@ def test_train_calls_train_pipeline_with_parsed_instructions( num_samples=500, num_epochs=3, train_datapoint_examples=None, - device=None + device=None, + train_dataset_path=None ) assert result is mock_output @@ -213,7 +214,8 @@ def test_train_calls_train_pipeline_with_all_arguments( num_samples=num_samples, num_epochs=num_epochs, train_datapoint_examples=train_examples, - device=None + device=None, + train_dataset_path=None ) assert result is mock_output @@ -332,7 +334,8 @@ def test_train_with_custom_num_samples( num_samples=custom_samples, num_epochs=3, train_datapoint_examples=None, - device=None + device=None, + train_dataset_path=None ) assert result is mock_output @@ -368,7 +371,8 @@ def test_train_with_custom_num_epochs( num_samples=500, num_epochs=custom_epochs, train_datapoint_examples=None, - device=None + device=None, + train_dataset_path=None ) assert result is mock_output @@ -404,7 +408,8 @@ def test_train_with_custom_output_path( num_samples=500, num_epochs=3, train_datapoint_examples=None, - device=None + device=None, + train_dataset_path=None ) assert result is mock_output @@ -443,7 +448,8 @@ def test_train_with_train_datapoint_examples( num_samples=500, num_epochs=3, train_datapoint_examples=examples, - device=None + device=None, + train_dataset_path=None ) assert result is mock_output diff --git a/tests/unit/spam_detection/test_sd__init__.py b/tests/unit/spam_detection/test_sd__init__.py index 402e6d7..c0a6cfb 100644 --- a/tests/unit/spam_detection/test_sd__init__.py +++ b/tests/unit/spam_detection/test_sd__init__.py @@ -61,6 +61,7 @@ def test_spam_detection_init_invalid_language_raises_validation_error(mocker: Mo ("english", "mocked-spam-english"), ("spanish", "mocked-spam-spanish"), ("german", "mocked-spam-german"), + ("italian", "mocked-spam-italian") ] ) def test_spam_detection_init_valid_language_uses_correct_model( @@ -81,6 +82,7 @@ def test_spam_detection_init_valid_language_uses_correct_model( mock_config.SPAM_DETECTION_ENGLISH_HF_BASE_MODEL = "mocked-spam-english" mock_config.SPAM_DETECTION_SPANISH_HF_BASE_MODEL = "mocked-spam-spanish" mock_config.SPAM_DETECTION_GERMAN_HF_BASE_MODEL = "mocked-spam-german" + mock_config.SPAM_DETECTION_ITALIAN_HF_BASE_MODEL = "mocked-spam-italian" # Mock ClassificationModel.__init__ mock_super_init = mocker.patch( "artifex.models.classification.classification_model.ClassificationModel.__init__", @@ -109,6 +111,7 @@ def test_spam_detection_init_default_language_is_english(mocker: MockerFixture): mock_config.SPAM_DETECTION_ENGLISH_HF_BASE_MODEL = "mocked-spam-english" mock_config.SPAM_DETECTION_SPANISH_HF_BASE_MODEL = "mocked-spam-spanish" mock_config.SPAM_DETECTION_GERMAN_HF_BASE_MODEL = "mocked-spam-german" + mock_config.SPAM_DETECTION_ITALIAN_HF_BASE_MODEL = "mocked-spam-italian" # Mock ClassificationModel.__init__ mock_super_init = mocker.patch( "artifex.models.classification.classification_model.ClassificationModel.__init__", diff --git a/tests/unit/text_summarization/__init__.py b/tests/unit/text_summarization/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/unit/text_summarization/test_ts__call__.py b/tests/unit/text_summarization/test_ts__call__.py new file mode 100644 index 0000000..2c6b3ec --- /dev/null +++ b/tests/unit/text_summarization/test_ts__call__.py @@ -0,0 +1,114 @@ +import pytest +from pytest_mock import MockerFixture +from synthex import Synthex +from unittest.mock import MagicMock +import torch + +from artifex.models.text_summarization import TextSummarization +from artifex.config import config + + +@pytest.fixture(scope="function", autouse=True) +def mock_dependencies(mocker: MockerFixture) -> None: + mocker.patch.object(config, "TEXT_SUMMARIZATION_HF_BASE_MODEL", "mock-summarization-model") + mocker.patch.object(config, "TEXT_SUMMARIZATION_MAX_TARGET_LENGTH", 128) + + # Mock AutoModelForSeq2SeqLM.from_pretrained + mock_model = MagicMock() + mocker.patch( + "artifex.models.text_summarization.text_summarization.AutoModelForSeq2SeqLM.from_pretrained", + return_value=mock_model + ) + + # Mock AutoTokenizer.from_pretrained + mock_tokenizer = MagicMock() + mock_tokenizer.return_value = { + "input_ids": torch.tensor([[1, 2, 3]]), + "attention_mask": torch.tensor([[1, 1, 1]]) + } + mocker.patch( + "artifex.models.text_summarization.text_summarization.AutoTokenizer.from_pretrained", + return_value=mock_tokenizer + ) + + +@pytest.fixture +def mock_synthex(mocker: MockerFixture) -> Synthex: + return mocker.MagicMock(spec=Synthex) + + +@pytest.fixture +def model(mock_synthex: Synthex) -> TextSummarization: + return TextSummarization(mock_synthex) + + +@pytest.mark.unit +def test_call_returns_list_of_strings( + model: TextSummarization, mocker: MockerFixture +) -> None: + """ + Test that __call__ returns a list of strings. + """ + # Mock model.generate to return some IDs + model._model.generate.return_value = torch.tensor([[1, 2, 3]]) + # Mock tokenizer.batch_decode to return some strings + model._tokenizer.batch_decode.return_value = ["A concise summary."] + + result = model(text="Some long text to summarize.", disable_logging=True) + assert isinstance(result, list) + assert all(isinstance(s, str) for s in result) + assert result == ["A concise summary."] + + +@pytest.mark.unit +def test_call_wraps_single_string_in_list( + model: TextSummarization, mocker: MockerFixture +) -> None: + """ + Test that __call__ handles a single string input (not a list). + """ + model._model.generate.return_value = torch.tensor([[1, 2, 3]]) + model._tokenizer.batch_decode.return_value = ["Summary."] + + model(text="A single text input.", disable_logging=True) + + # Check that tokenizer was called with a list + call_args = model._tokenizer.call_args[0][0] + assert isinstance(call_args, list) + assert len(call_args) == 1 + assert call_args[0] == "A single text input." + + +@pytest.mark.unit +def test_call_handles_list_of_texts( + model: TextSummarization, mocker: MockerFixture +) -> None: + """ + Test that __call__ handles a list of texts and returns one summary per input. + """ + texts = ["First long text.", "Second long text.", "Third long text."] + model._model.generate.return_value = torch.tensor([[1, 2, 3]] * 3) + model._tokenizer.batch_decode.return_value = ["Summary 0.", "Summary 1.", "Summary 2."] + + result = model(text=texts, disable_logging=True) + assert len(result) == len(texts) + assert result == ["Summary 0.", "Summary 1.", "Summary 2."] + + +@pytest.mark.unit +def test_call_uses_correct_parameters( + model: TextSummarization, mocker: MockerFixture +) -> None: + """ + Test that __call__ uses correct parameters for generation. + """ + model._model.generate.return_value = torch.tensor([[1, 2, 3]]) + model._tokenizer.batch_decode.return_value = ["Summary."] + + model(text="Some text.", disable_logging=True) + + # Check model.generate call + model._model.generate.assert_called_once() + call_kwargs = model._model.generate.call_args.kwargs + assert call_kwargs["max_length"] == config.TEXT_SUMMARIZATION_MAX_TARGET_LENGTH + assert call_kwargs["num_beams"] == 4 diff --git a/tests/unit/text_summarization/test_ts_get_data_gen_instr.py b/tests/unit/text_summarization/test_ts_get_data_gen_instr.py new file mode 100644 index 0000000..e14aafc --- /dev/null +++ b/tests/unit/text_summarization/test_ts_get_data_gen_instr.py @@ -0,0 +1,93 @@ +import pytest +from pytest_mock import MockerFixture +from synthex import Synthex +from unittest.mock import MagicMock + +from artifex.models.text_summarization import TextSummarization +from artifex.core import ParsedModelInstructions + + +@pytest.fixture +def mock_dependencies(mocker: MockerFixture) -> None: + mocker.patch( + "artifex.models.text_summarization.text_summarization.AutoModelForSeq2SeqLM.from_pretrained", + return_value=MagicMock() + ) + mocker.patch( + "artifex.models.text_summarization.text_summarization.AutoTokenizer.from_pretrained", + return_value=MagicMock() + ) + + +@pytest.fixture +def mock_synthex(mocker: MockerFixture) -> Synthex: + return mocker.MagicMock() + + +@pytest.fixture +def model(mock_dependencies: None, mock_synthex: Synthex) -> TextSummarization: + return TextSummarization(synthex=mock_synthex) + + +@pytest.mark.unit +def test_get_data_gen_instr_returns_list_of_strings( + model: TextSummarization, +) -> None: + """ + Test that _get_data_gen_instr returns a list of strings. + """ + user_instr = ParsedModelInstructions( + user_instructions=["medical research"], + language="english" + ) + result = model._get_data_gen_instr(user_instr) + assert isinstance(result, list) + assert all(isinstance(item, str) for item in result) + + +@pytest.mark.unit +def test_get_data_gen_instr_formats_language_placeholder( + model: TextSummarization, +) -> None: + """ + Test that _get_data_gen_instr correctly formats the language placeholder. + """ + user_instr = ParsedModelInstructions( + user_instructions=["science"], + language="french" + ) + result = model._get_data_gen_instr(user_instr) + language_instructions = [instr for instr in result if "french" in instr.lower()] + assert len(language_instructions) > 0 + + +@pytest.mark.unit +def test_get_data_gen_instr_formats_domain_placeholder( + model: TextSummarization, +) -> None: + """ + Test that _get_data_gen_instr correctly formats the domain placeholder. + """ + domain = "technology news" + user_instr = ParsedModelInstructions( + user_instructions=[domain], + language="english" + ) + result = model._get_data_gen_instr(user_instr) + domain_instructions = [instr for instr in result if domain in instr] + assert len(domain_instructions) > 0 + + +@pytest.mark.unit +def test_get_data_gen_instr_length_matches_system_instructions( + model: TextSummarization, +) -> None: + """ + Test that the number of returned instructions equals the number of system instructions. + """ + user_instr = ParsedModelInstructions( + user_instructions=["sports"], + language="english" + ) + result = model._get_data_gen_instr(user_instr) + assert len(result) == len(model._system_data_gen_instr) diff --git a/tests/unit/text_summarization/test_ts_load_model.py b/tests/unit/text_summarization/test_ts_load_model.py new file mode 100644 index 0000000..bd2dfdc --- /dev/null +++ b/tests/unit/text_summarization/test_ts_load_model.py @@ -0,0 +1,77 @@ +import pytest +from pytest_mock import MockerFixture +from typing import Any + +from artifex.models.text_summarization import TextSummarization + + +@pytest.fixture +def mock_synthex(mocker: MockerFixture) -> Any: + return mocker.Mock() + + +@pytest.fixture +def model(mock_synthex: Any, mocker: MockerFixture) -> TextSummarization: + mocker.patch( + "artifex.models.text_summarization.text_summarization.AutoModelForSeq2SeqLM.from_pretrained" + ) + mocker.patch( + "artifex.models.text_summarization.text_summarization.AutoTokenizer.from_pretrained" + ) + mock_config = mocker.patch("artifex.models.text_summarization.text_summarization.config") + mock_config.TEXT_SUMMARIZATION_HF_BASE_MODEL = "mock-model" + mock_config.DEFAULT_SYNTHEX_DATAPOINT_NUM = 100 + return TextSummarization(mock_synthex) + + +@pytest.mark.unit +def test_load_model_calls_seq2seq_from_pretrained( + model: TextSummarization, mocker: MockerFixture +) -> None: + """ + Test that _load_model calls AutoModelForSeq2SeqLM.from_pretrained with the model path. + """ + mock_from_pretrained = mocker.patch( + "artifex.models.text_summarization.text_summarization.AutoModelForSeq2SeqLM.from_pretrained" + ) + mocker.patch( + "artifex.models.text_summarization.text_summarization.AutoTokenizer.from_pretrained" + ) + model._load_model("/path/to/model") + mock_from_pretrained.assert_called_once_with("/path/to/model") + + +@pytest.mark.unit +def test_load_model_calls_tokenizer_from_pretrained( + model: TextSummarization, mocker: MockerFixture +) -> None: + """ + Test that _load_model also reloads the tokenizer from the model path. + """ + mocker.patch( + "artifex.models.text_summarization.text_summarization.AutoModelForSeq2SeqLM.from_pretrained" + ) + mock_tokenizer_from_pretrained = mocker.patch( + "artifex.models.text_summarization.text_summarization.AutoTokenizer.from_pretrained" + ) + model._load_model("/path/to/model") + mock_tokenizer_from_pretrained.assert_called_once_with("/path/to/model") + + +@pytest.mark.unit +def test_load_model_updates_model_instance( + model: TextSummarization, mocker: MockerFixture +) -> None: + """ + Test that _load_model updates the _model instance variable. + """ + new_model = mocker.MagicMock() + mocker.patch( + "artifex.models.text_summarization.text_summarization.AutoModelForSeq2SeqLM.from_pretrained", + return_value=new_model + ) + mocker.patch( + "artifex.models.text_summarization.text_summarization.AutoTokenizer.from_pretrained" + ) + model._load_model("/path/to/model") + assert model._model is new_model diff --git a/tests/unit/text_summarization/test_ts_parse_user_instructions.py b/tests/unit/text_summarization/test_ts_parse_user_instructions.py new file mode 100644 index 0000000..44b93b5 --- /dev/null +++ b/tests/unit/text_summarization/test_ts_parse_user_instructions.py @@ -0,0 +1,64 @@ +import pytest +from pytest_mock import MockerFixture +from synthex import Synthex +from unittest.mock import MagicMock + +from artifex.models.text_summarization import TextSummarization +from artifex.core import ParsedModelInstructions +from artifex.config import config + + +@pytest.fixture +def mock_dependencies(mocker: MockerFixture) -> None: + mocker.patch( + "artifex.models.text_summarization.text_summarization.AutoModelForSeq2SeqLM.from_pretrained", + return_value=MagicMock() + ) + mocker.patch( + "artifex.models.text_summarization.text_summarization.AutoTokenizer.from_pretrained", + return_value=MagicMock() + ) + + +@pytest.fixture +def mock_synthex(mocker: MockerFixture) -> Synthex: + return mocker.MagicMock() + + +@pytest.fixture +def model(mock_dependencies: None, mock_synthex: Synthex) -> TextSummarization: + return TextSummarization(synthex=mock_synthex) + + +@pytest.mark.unit +def test_parse_user_instructions_returns_parsed_model_instructions( + model: TextSummarization, +) -> None: + """ + Test that _parse_user_instructions returns a ParsedModelInstructions instance. + """ + result = model._parse_user_instructions("medical research", "english") + assert isinstance(result, ParsedModelInstructions) + + +@pytest.mark.unit +def test_parse_user_instructions_stores_domain_in_user_instructions( + model: TextSummarization, +) -> None: + """ + Test that the domain string ends up in the user_instructions list. + """ + domain = "financial news" + result = model._parse_user_instructions(domain, "english") + assert domain in result.user_instructions + + +@pytest.mark.unit +def test_parse_user_instructions_stores_language( + model: TextSummarization, +) -> None: + """ + Test that the language is stored correctly. + """ + result = model._parse_user_instructions("science", "spanish") + assert result.language == "spanish" diff --git a/tests/unit/text_summarization/test_ts_perform_train_pipeline.py b/tests/unit/text_summarization/test_ts_perform_train_pipeline.py new file mode 100644 index 0000000..3bd871d --- /dev/null +++ b/tests/unit/text_summarization/test_ts_perform_train_pipeline.py @@ -0,0 +1,153 @@ +import pytest +from pytest_mock import MockerFixture +from synthex import Synthex +from unittest.mock import MagicMock +from transformers.trainer_utils import TrainOutput +from datasets import DatasetDict, Dataset + +from artifex.models.text_summarization import TextSummarization +from artifex.core import ParsedModelInstructions +from artifex.config import config + + +@pytest.fixture +def mock_dependencies(mocker: MockerFixture) -> None: + mocker.patch( + "artifex.models.text_summarization.text_summarization.AutoModelForSeq2SeqLM.from_pretrained", + return_value=MagicMock() + ) + mocker.patch( + "artifex.models.text_summarization.text_summarization.AutoTokenizer.from_pretrained", + return_value=MagicMock() + ) + mocker.patch( + "artifex.models.text_summarization.text_summarization.torch.cuda.is_available", + return_value=False + ) + mocker.patch( + "artifex.models.text_summarization.text_summarization.torch.backends.mps.is_available", + return_value=False + ) + + +@pytest.fixture +def mock_synthex(mocker: MockerFixture) -> Synthex: + return mocker.MagicMock() + + +@pytest.fixture +def mock_get_model_output_path(mocker: MockerFixture) -> MagicMock: + return mocker.patch( + "artifex.models.text_summarization.text_summarization.get_model_output_path", + return_value="/test/output/model" + ) + + +@pytest.fixture +def mock_seq2seq_training_args(mocker: MockerFixture) -> MagicMock: + return mocker.patch( + "artifex.models.text_summarization.text_summarization.Seq2SeqTrainingArguments" + ) + + +@pytest.fixture +def mock_silent_seq2seq_trainer(mocker: MockerFixture) -> MagicMock: + mock_trainer_instance = MagicMock() + mock_trainer_instance.train.return_value = TrainOutput( + global_step=100, training_loss=0.5, metrics={} + ) + mock_trainer_class = mocker.patch( + "artifex.models.text_summarization.text_summarization.SilentSeq2SeqTrainer", + return_value=mock_trainer_instance + ) + return mock_trainer_class + + +@pytest.fixture +def model(mock_dependencies: None, mock_synthex: Synthex) -> TextSummarization: + return TextSummarization(synthex=mock_synthex) + + +@pytest.mark.unit +def test_perform_train_pipeline_calls_silent_seq2seq_trainer( + model: TextSummarization, + mocker: MockerFixture, + mock_get_model_output_path: MagicMock, + mock_seq2seq_training_args: MagicMock, + mock_silent_seq2seq_trainer: MagicMock, +) -> None: + """ + Test that _perform_train_pipeline instantiates SilentSeq2SeqTrainer. + """ + mock_dataset = MagicMock(spec=DatasetDict) + mock_dataset.__getitem__ = MagicMock(return_value=MagicMock(spec=Dataset)) + mocker.patch.object(model, "_build_tokenized_train_ds", return_value=mock_dataset) + + user_instr = ParsedModelInstructions( + user_instructions=["medical research"], + language="english" + ) + model._perform_train_pipeline( + user_instructions=user_instr, + output_path="/test/output/", + num_samples=50, + num_epochs=1, + ) + mock_silent_seq2seq_trainer.assert_called_once() + + +@pytest.mark.unit +def test_perform_train_pipeline_returns_train_output( + model: TextSummarization, + mocker: MockerFixture, + mock_get_model_output_path: MagicMock, + mock_seq2seq_training_args: MagicMock, + mock_silent_seq2seq_trainer: MagicMock, +) -> None: + """ + Test that _perform_train_pipeline returns a TrainOutput instance. + """ + mock_dataset = MagicMock(spec=DatasetDict) + mock_dataset.__getitem__ = MagicMock(return_value=MagicMock(spec=Dataset)) + mocker.patch.object(model, "_build_tokenized_train_ds", return_value=mock_dataset) + + user_instr = ParsedModelInstructions( + user_instructions=["technology"], + language="english" + ) + result = model._perform_train_pipeline( + user_instructions=user_instr, + output_path="/test/output/", + num_samples=50, + num_epochs=1, + ) + assert isinstance(result, TrainOutput) + + +@pytest.mark.unit +def test_perform_train_pipeline_saves_model( + model: TextSummarization, + mocker: MockerFixture, + mock_get_model_output_path: MagicMock, + mock_seq2seq_training_args: MagicMock, + mock_silent_seq2seq_trainer: MagicMock, +) -> None: + """ + Test that trainer.save_model() is called after training. + """ + mock_dataset = MagicMock(spec=DatasetDict) + mock_dataset.__getitem__ = MagicMock(return_value=MagicMock(spec=Dataset)) + mocker.patch.object(model, "_build_tokenized_train_ds", return_value=mock_dataset) + + user_instr = ParsedModelInstructions( + user_instructions=["law"], + language="english" + ) + model._perform_train_pipeline( + user_instructions=user_instr, + output_path="/test/output/", + num_samples=50, + num_epochs=1, + ) + trainer_instance = mock_silent_seq2seq_trainer.return_value + trainer_instance.save_model.assert_called_once() diff --git a/tests/unit/text_summarization/test_ts_post_process_synthetic_dataset.py b/tests/unit/text_summarization/test_ts_post_process_synthetic_dataset.py new file mode 100644 index 0000000..8de0453 --- /dev/null +++ b/tests/unit/text_summarization/test_ts_post_process_synthetic_dataset.py @@ -0,0 +1,148 @@ +import pytest +from pytest_mock import MockerFixture +from synthex import Synthex +from unittest.mock import MagicMock +import pandas as pd +import tempfile +import os + +from artifex.models.text_summarization import TextSummarization + + +@pytest.fixture +def mock_dependencies(mocker: MockerFixture) -> None: + mocker.patch( + "artifex.models.text_summarization.text_summarization.AutoModelForSeq2SeqLM.from_pretrained", + return_value=MagicMock() + ) + mocker.patch( + "artifex.models.text_summarization.text_summarization.AutoTokenizer.from_pretrained", + return_value=MagicMock() + ) + + +@pytest.fixture +def mock_synthex(mocker: MockerFixture) -> Synthex: + return mocker.MagicMock() + + +@pytest.fixture +def model(mock_dependencies: None, mock_synthex: Synthex) -> TextSummarization: + return TextSummarization(synthex=mock_synthex) + + +@pytest.fixture +def temp_csv_file(): + temp_file = tempfile.NamedTemporaryFile(mode="w", suffix=".csv", delete=False) + temp_path = temp_file.name + temp_file.close() + yield temp_path + if os.path.exists(temp_path): + os.remove(temp_path) + + +def _long_text(n: int = 100) -> str: + return "word " * n + + +@pytest.mark.unit +def test_removes_rows_with_null_text( + model: TextSummarization, temp_csv_file: str +) -> None: + """ + Test that rows with null text are removed. + """ + df = pd.DataFrame({ + "text": [None, _long_text()], + "summary": ["A short summary.", "Another short summary."], + }) + df.to_csv(temp_csv_file, index=False) + model._post_process_synthetic_dataset(temp_csv_file) + result = pd.read_csv(temp_csv_file) + assert len(result) == 1 + + +@pytest.mark.unit +def test_removes_rows_with_short_text( + model: TextSummarization, temp_csv_file: str +) -> None: + """ + Test that rows where text is shorter than 50 characters are removed. + """ + df = pd.DataFrame({ + "text": ["Short.", _long_text()], + "summary": ["Summary one.", "Summary two."], + }) + df.to_csv(temp_csv_file, index=False) + model._post_process_synthetic_dataset(temp_csv_file) + result = pd.read_csv(temp_csv_file) + assert len(result) == 1 + + +@pytest.mark.unit +def test_removes_rows_with_null_summary( + model: TextSummarization, temp_csv_file: str +) -> None: + """ + Test that rows with null summaries are removed. + """ + df = pd.DataFrame({ + "text": [_long_text(), _long_text()], + "summary": [None, "Valid summary here."], + }) + df.to_csv(temp_csv_file, index=False) + model._post_process_synthetic_dataset(temp_csv_file) + result = pd.read_csv(temp_csv_file) + assert len(result) == 1 + + +@pytest.mark.unit +def test_removes_rows_with_short_summary( + model: TextSummarization, temp_csv_file: str +) -> None: + """ + Test that rows where summary is shorter than 10 characters are removed. + """ + df = pd.DataFrame({ + "text": [_long_text(), _long_text()], + "summary": ["Too short", "This is a valid summary sentence."], + }) + df.to_csv(temp_csv_file, index=False) + model._post_process_synthetic_dataset(temp_csv_file) + result = pd.read_csv(temp_csv_file) + assert len(result) == 1 + + +@pytest.mark.unit +def test_removes_rows_where_summary_is_not_shorter_than_text( + model: TextSummarization, temp_csv_file: str +) -> None: + """ + Test that rows where the summary is as long or longer than the text are removed. + """ + long_text = _long_text(50) + df = pd.DataFrame({ + "text": [long_text, long_text], + "summary": [long_text, "A brief summary."], + }) + df.to_csv(temp_csv_file, index=False) + model._post_process_synthetic_dataset(temp_csv_file) + result = pd.read_csv(temp_csv_file) + assert len(result) == 1 + + +@pytest.mark.unit +def test_valid_rows_are_kept( + model: TextSummarization, temp_csv_file: str +) -> None: + """ + Test that rows with valid text and summary are retained. + """ + df = pd.DataFrame({ + "text": [_long_text() for _ in range(5)], + "summary": ["Valid summary sentence here." for _ in range(5)], + }) + df.to_csv(temp_csv_file, index=False) + model._post_process_synthetic_dataset(temp_csv_file) + result = pd.read_csv(temp_csv_file) + assert len(result) == 5 diff --git a/tests/unit/text_summarization/test_ts_synthetic_to_training_dataset.py b/tests/unit/text_summarization/test_ts_synthetic_to_training_dataset.py new file mode 100644 index 0000000..1d7bfe9 --- /dev/null +++ b/tests/unit/text_summarization/test_ts_synthetic_to_training_dataset.py @@ -0,0 +1,88 @@ +import pytest +from pytest_mock import MockerFixture +from synthex import Synthex +from unittest.mock import MagicMock +from datasets import DatasetDict +import pandas as pd +import tempfile +import os + +from artifex.models.text_summarization import TextSummarization + + +@pytest.fixture +def mock_dependencies(mocker: MockerFixture) -> None: + mocker.patch( + "artifex.models.text_summarization.text_summarization.AutoModelForSeq2SeqLM.from_pretrained", + return_value=MagicMock() + ) + mocker.patch( + "artifex.models.text_summarization.text_summarization.AutoTokenizer.from_pretrained", + return_value=MagicMock() + ) + + +@pytest.fixture +def mock_synthex(mocker: MockerFixture) -> Synthex: + return mocker.MagicMock() + + +@pytest.fixture +def model(mock_dependencies: None, mock_synthex: Synthex) -> TextSummarization: + return TextSummarization(synthex=mock_synthex) + + +@pytest.fixture +def temp_csv_file(): + temp_file = tempfile.NamedTemporaryFile(mode="w", suffix=".csv", delete=False) + temp_path = temp_file.name + temp_file.close() + yield temp_path + if os.path.exists(temp_path): + os.remove(temp_path) + + +def _make_csv(temp_csv_file: str, n: int = 10) -> None: + df = pd.DataFrame({ + "text": [f"This is a longer text that is sentence number {i}. " * 5 for i in range(n)], + "summary": [f"Summary {i}." for i in range(n)], + }) + df.to_csv(temp_csv_file, index=False) + + +@pytest.mark.unit +def test_returns_dataset_dict( + model: TextSummarization, temp_csv_file: str +) -> None: + """ + Test that _synthetic_to_training_dataset returns a DatasetDict. + """ + _make_csv(temp_csv_file) + result = model._synthetic_to_training_dataset(temp_csv_file) + assert isinstance(result, DatasetDict) + + +@pytest.mark.unit +def test_dataset_has_train_and_test_splits( + model: TextSummarization, temp_csv_file: str +) -> None: + """ + Test that the returned DatasetDict has 'train' and 'test' keys. + """ + _make_csv(temp_csv_file) + result = model._synthetic_to_training_dataset(temp_csv_file) + assert "train" in result + assert "test" in result + + +@pytest.mark.unit +def test_dataset_contains_text_and_summary_columns( + model: TextSummarization, temp_csv_file: str +) -> None: + """ + Test that both 'text' and 'summary' columns are present in the dataset. + """ + _make_csv(temp_csv_file) + result = model._synthetic_to_training_dataset(temp_csv_file) + assert "text" in result["train"].column_names + assert "summary" in result["train"].column_names diff --git a/tests/unit/text_summarization/test_ts_tokenize_dataset.py b/tests/unit/text_summarization/test_ts_tokenize_dataset.py new file mode 100644 index 0000000..078c4e4 --- /dev/null +++ b/tests/unit/text_summarization/test_ts_tokenize_dataset.py @@ -0,0 +1,108 @@ +import pytest +from pytest_mock import MockerFixture +from synthex import Synthex +from unittest.mock import MagicMock +from datasets import DatasetDict, Dataset + +from artifex.models.text_summarization import TextSummarization +from artifex.config import config + + +@pytest.fixture +def mock_dependencies(mocker: MockerFixture) -> None: + mocker.patch( + "artifex.models.text_summarization.text_summarization.AutoModelForSeq2SeqLM.from_pretrained", + return_value=MagicMock() + ) + mocker.patch( + "artifex.models.text_summarization.text_summarization.AutoTokenizer.from_pretrained", + return_value=MagicMock() + ) + + +@pytest.fixture +def mock_synthex(mocker: MockerFixture) -> Synthex: + return mocker.MagicMock() + + +@pytest.fixture +def model(mock_dependencies: None, mock_synthex: Synthex) -> TextSummarization: + return TextSummarization(synthex=mock_synthex) + + +def _make_dataset() -> DatasetDict: + data = { + "text": ["This is a long example text for testing purposes."] * 10, + "summary": ["Short summary."] * 10, + } + dataset = Dataset.from_dict(data) + return dataset.train_test_split(test_size=0.2) + + +@pytest.mark.unit +def test_tokenize_dataset_returns_dataset_dict( + model: TextSummarization, mocker: MockerFixture +) -> None: + """ + Test that _tokenize_dataset returns a DatasetDict. + """ + + def fake_tokenizer(texts=None, text_target=None, **kwargs): + n = len(texts) if texts is not None else len(text_target) + return {"input_ids": [[1, 2, 3]] * n, "attention_mask": [[1, 1, 1]] * n} + + mock_tokenizer = MagicMock(side_effect=fake_tokenizer) + mock_tokenizer.pad_token_id = 0 + model._tokenizer_val = mock_tokenizer + + dataset = _make_dataset() + result = model._tokenize_dataset(dataset, model._token_keys) + assert isinstance(result, DatasetDict) + + +@pytest.mark.unit +def test_tokenize_dataset_produces_labels_column( + model: TextSummarization, mocker: MockerFixture +) -> None: + """ + Test that _tokenize_dataset adds a 'labels' column to the dataset. + """ + + def fake_tokenizer(texts=None, text_target=None, **kwargs): + n = len(texts) if texts is not None else len(text_target) + return {"input_ids": [[1, 2, 3]] * n, "attention_mask": [[1, 1, 1]] * n} + + mock_tokenizer = MagicMock(side_effect=fake_tokenizer) + mock_tokenizer.pad_token_id = 0 + model._tokenizer_val = mock_tokenizer + + dataset = _make_dataset() + result = model._tokenize_dataset(dataset, model._token_keys) + assert "labels" in result["train"].column_names + + +@pytest.mark.unit +def test_tokenize_dataset_replaces_pad_tokens_with_minus_100( + model: TextSummarization, mocker: MockerFixture +) -> None: + """ + Test that padding token IDs in labels are replaced with -100. + """ + pad_id = 0 + + def fake_tokenizer(texts=None, text_target=None, **kwargs): + n = len(texts) if texts is not None else len(text_target) + if text_target is not None: + return {"input_ids": [[1, pad_id, 3]] * n, "attention_mask": [[1, 1, 1]] * n} + return {"input_ids": [[1, 2, 3]] * n, "attention_mask": [[1, 1, 1]] * n} + + mock_tokenizer = MagicMock(side_effect=fake_tokenizer) + mock_tokenizer.pad_token_id = pad_id + model._tokenizer_val = mock_tokenizer + + dataset = _make_dataset() + result = model._tokenize_dataset(dataset, model._token_keys) + + for label_seq in result["train"]["labels"]: + assert pad_id not in label_seq + assert -100 in label_seq diff --git a/tests/unit/text_summarization/test_ts_train.py b/tests/unit/text_summarization/test_ts_train.py new file mode 100644 index 0000000..8bb9176 --- /dev/null +++ b/tests/unit/text_summarization/test_ts_train.py @@ -0,0 +1,93 @@ +import pytest +from pytest_mock import MockerFixture +from synthex import Synthex +from unittest.mock import MagicMock +from transformers.trainer_utils import TrainOutput + +from artifex.models.text_summarization import TextSummarization +from artifex.config import config + + +@pytest.fixture(scope="function", autouse=True) +def mock_dependencies(mocker: MockerFixture) -> None: + mocker.patch.object(config, "TEXT_SUMMARIZATION_HF_BASE_MODEL", "mock-summarization-model") + mocker.patch( + "artifex.models.text_summarization.text_summarization.AutoModelForSeq2SeqLM.from_pretrained", + return_value=MagicMock() + ) + mocker.patch( + "artifex.models.text_summarization.text_summarization.AutoTokenizer.from_pretrained", + return_value=MagicMock() + ) + + +@pytest.fixture +def mock_synthex(mocker: MockerFixture) -> Synthex: + return mocker.MagicMock(spec=Synthex) + + +@pytest.fixture +def model(mock_synthex: Synthex) -> TextSummarization: + return TextSummarization(mock_synthex) + + +@pytest.mark.unit +def test_train_calls_train_pipeline( + model: TextSummarization, mocker: MockerFixture +) -> None: + """ + Test that train() calls _train_pipeline. + """ + mocker.patch.object( + model, "_train_pipeline", + return_value=TrainOutput(global_step=1, training_loss=0.1, metrics={}) + ) + model.train(domain="news articles", language="english", disable_logging=True) + model._train_pipeline.assert_called_once() + + +@pytest.mark.unit +def test_train_passes_correct_num_samples( + model: TextSummarization, mocker: MockerFixture +) -> None: + """ + Test that train() forwards num_samples to _train_pipeline. + """ + mock_pipeline = mocker.patch.object( + model, "_train_pipeline", + return_value=TrainOutput(global_step=1, training_loss=0.1, metrics={}) + ) + model.train(domain="science", num_samples=200, disable_logging=True) + _, kwargs = mock_pipeline.call_args + assert kwargs.get("num_samples") == 200 + + +@pytest.mark.unit +def test_train_passes_correct_num_epochs( + model: TextSummarization, mocker: MockerFixture +) -> None: + """ + Test that train() forwards num_epochs to _train_pipeline. + """ + mock_pipeline = mocker.patch.object( + model, "_train_pipeline", + return_value=TrainOutput(global_step=1, training_loss=0.1, metrics={}) + ) + model.train(domain="finance", num_epochs=5, disable_logging=True) + _, kwargs = mock_pipeline.call_args + assert kwargs.get("num_epochs") == 5 + + +@pytest.mark.unit +def test_train_returns_train_output( + model: TextSummarization, mocker: MockerFixture +) -> None: + """ + Test that train() returns a TrainOutput instance. + """ + mocker.patch.object( + model, "_train_pipeline", + return_value=TrainOutput(global_step=1, training_loss=0.1, metrics={}) + ) + result = model.train(domain="history", disable_logging=True) + assert isinstance(result, TrainOutput)