standard.py

from typing import List, Optional, Union

from .augmentors import (
 Augmentor,
 FinalStateInputsAugmentor,
 NullAugmentor,
 TaskInputsAugmentor,
)
from .card import TaskCard
from .collections_operators import GetLength
from .dataclass import Field, InternalField, NonPositionalField, OptionalField
from .formats import Format, SystemFormat
from .logging_utils import get_logger
from .operator import SequentialOperator, SourceSequentialOperator, StreamingOperator
from .operators import Set, StreamRefiner
from .recipe import Recipe
from .schema import Finalize
from .serializers import SingleTypeSerializer
from .settings_utils import get_constants
from .splitters import ConstantSizeSample, RandomSizeSample, Sampler, SeparateSplit
from .stream import MultiStream
from .system_prompts import EmptySystemPrompt, SystemPrompt
from .task import Task
from .templates import ApplyRandomTemplate, ApplySingleTemplate, Template, TemplatesList

constants = get_constants()
logger = get_logger()


# Used to give meaningful name to recipe steps
class CreateDemosPool(SeparateSplit):
 pass


class BaseRecipe(Recipe, SourceSequentialOperator):
 # Base parameters
 card: TaskCard = None
 task: Task = None
 template: Union[Template, List[Template], TemplatesList] = None
 system_prompt: SystemPrompt = Field(default_factory=EmptySystemPrompt)
 format: Format = Field(default_factory=SystemFormat)
 serializer: Union[SingleTypeSerializer, List[SingleTypeSerializer]] = None

 # Additional parameters
 template_card_index: int = NonPositionalField(default=None)
 metrics: List[str] = NonPositionalField(default=None)
 postprocessors: List[str] = NonPositionalField(default=None)

 group_by: List[Union[str, List[str]]] = []

 loader_limit: int = None

 max_train_instances: int = None
 max_validation_instances: int = None
 max_test_instances: int = None

 train_refiner: StreamRefiner = OptionalField(default_factory=StreamRefiner)
 validation_refiner: StreamRefiner = OptionalField(default_factory=StreamRefiner)
 test_refiner: StreamRefiner = OptionalField(default_factory=StreamRefiner)

 demos_pool_size: int = None
 num_demos: Optional[Union[int, List[int]]] = 0
 demos_removed_from_data: bool = True

 demos_pool_name: str = "demos_pool"
 demos_taken_from: str = "train"
 demos_field: str = "demos"
 sampler: Sampler = None

 augmentor: Augmentor = OptionalField(default_factory=NullAugmentor)

 steps: List[StreamingOperator] = InternalField(default_factory=list)

 def before_process_multi_stream(self):
 super().before_process_multi_stream()

 @property
 def max_demos_size(self):
 if isinstance(self.num_demos, list):
 return max(self.num_demos)
 return self.num_demos

 def verify(self):
 super().verify()

 if self.task is None and self.card is None:
 raise ValueError("Set card or task in the recipe")

 if self.card is None and (
 self.num_demos > 0 or self.demos_pool_size is not None
 ):
 raise ValueError(
 "To use num_demos and demos_pool_size in recipe set a card."
 )

 if self.use_demos:
 if self.demos_pool_size is None or self.demos_pool_size < 1:
 raise ValueError(
 "When using demonstrations both num_demos and demos_pool_size should be assigned with positive integers."
 )
 if self.demos_pool_size < self.max_demos_size:
 raise ValueError(
 f"num_demos (got: {self.max_demos_size}) should not exceed demos_pool_size (got: {self.demos_pool_size})"
 )
 if self.loader_limit and self.demos_pool_size > self.loader_limit:
 raise ValueError(
 f"demos_pool_size should not exceed loader_limit ({self.loader_limit}), Got demos_pool_size={self.demos_pool_size}"
 )

 if self.loader_limit:
 if self.max_test_instances and self.max_test_instances > self.loader_limit:
 raise ValueError(
 f"max_test_instances should not exceed loader_limit ({self.loader_limit}), Got max_test_instances={self.max_test_instances}"
 )
 if (
 self.max_validation_instances
 and self.max_validation_instances > self.loader_limit
 ):
 raise ValueError(
 f"max_validation_instances should not exceed loader_limit ({self.loader_limit}), Got max_validation_instances={self.max_validation_instances}"
 )
 if (
 self.max_train_instances
 and self.max_train_instances > self.loader_limit
 ):
 raise ValueError(
 f"max_train_instances should not exceed loader_limit ({self.loader_limit}), Got max_train_instances={self.max_train_instances}"
 )
 if self.metrics is not None and not isinstance(self.metrics, List):
 raise ValueError(
 f"metrics must be a list of metrics. Got metrics = {self.metrics}"
 )
 if self.postprocessors is not None and not isinstance(
 self.postprocessors, List
 ):
 raise ValueError(
 f"post processors must be a list of post processor. Got postprocessors = {self.postprocessors}"
 )

 if self.template is None:
 raise ValueError(
 "You must set in the recipe either `template`, `template_card_index` or `templates`."
 )

 if isinstance(self.template, list):
 for template in self.template:
 self.verify_template(template)
 else:
 self.verify_template(self.template)

 if self.serializer is not None:
 if not isinstance(self.serializer, list):
 self.serializer = [self.serializer]
 self.template.serializer.add_serializers(self.serializer)

 def prepare_refiners(self):
 self.train_refiner.max_instances = self.max_train_instances
 self.train_refiner.apply_to_streams = ["train"]
 self.processing.steps.append(self.train_refiner)

 self.validation_refiner.max_instances = self.max_validation_instances
 self.validation_refiner.apply_to_streams = ["validation"]
 self.processing.steps.append(self.validation_refiner)

 self.test_refiner.max_instances = self.max_test_instances
 self.test_refiner.apply_to_streams = ["test"]
 self.processing.steps.append(self.test_refiner)

 def verify_template(self, template):
 if not isinstance(template, Template):
 raise ValueError(
 f"template argument must be an object of type Template. Got template = {template}"
 )

 def set_pipelines(self):
 self.loading = SequentialOperator(
 __description__="Loading the data from the data source."
 )
 self.metadata = SequentialOperator(
 __description__="Adding metadata (e.g. format, system prompt, template) "
 )
 self.standardization = SequentialOperator(
 __description__="Standardizing the raw dataset fields to task field definition."
 )

 self.processing = SequentialOperator(
 __description__="Setting task fields (and selecting demos per sample if needed)."
 )
 self.verbalization = SequentialOperator()
 self.verbalization.__description__ = "Verbalizing the input to the model and gold references to the 'source', 'target' and 'references' fields."
 self.finalize = SequentialOperator()
 self.finalize.__description__ = "Adding post processors. Removing intermediate fields. Creating the final output dataset."

 self.steps = [
 self.loading,
 self.metadata,
 self.standardization,
 self.processing,
 self.metadata,
 self.verbalization,
 self.finalize,
 ]

 self.inference_instance = SequentialOperator()

 self.inference_instance.steps = [
 self.metadata,
 self.processing,
 self.metadata,
 ]

 self.inference_demos = SourceSequentialOperator()

 self.inference_demos.steps = [
 self.loading,
 self.metadata,
 self.standardization,
 self.processing,
 self.metadata,
 ]

 self.inference = SequentialOperator()

 self.inference.steps = [self.verbalization, self.finalize]

 self._demos_pool_cache = None

 def production_preprocess(self, task_instances):
 ms = MultiStream.from_iterables({constants.inference_stream: task_instances})
 return list(self.inference_instance(ms)[constants.inference_stream])

 def production_demos_pool(self):
 if self.use_demos:
 if self._demos_pool_cache is None:
 self._demos_pool_cache = list(
 self.inference_demos()[self.demos_pool_name]
 )
 return self._demos_pool_cache
 return []

 @property
 def has_custom_demos_pool(self):
 return self.demos_pool_size is not None and self.demos_pool_size > 0

 @property
 def use_demos(self):
 return self.num_demos is not None and self.max_demos_size > 0

 def produce(self, task_instances):
 """Use the recipe in production to produce model ready query from standard task instance."""
 self.before_process_multi_stream()
 streams = {
 constants.inference_stream: self.production_preprocess(task_instances),
 }
 if self.use_demos:
 streams[self.demos_pool_name] = self.production_demos_pool()
 multi_stream = MultiStream.from_iterables(streams)
 multi_stream = self.inference(multi_stream)
 return list(multi_stream[constants.inference_stream])

 def reset_pipeline(self):
 if self.card and self.card.preprocess_steps is None:
 self.card.preprocess_steps = []

 if self.task is None:
 self.task = self.card.task

 self.set_pipelines()

 if self.card is not None:
 loader = self.card.loader
 if self.loader_limit:
 loader.loader_limit = self.loader_limit
 logger.info(f"Loader line limit was set to {self.loader_limit}")
 self.loading.steps.append(loader)

 # This is required in case loader_limit is not enforced by the loader
 if self.loader_limit:
 self.loading.steps.append(
 StreamRefiner(max_instances=self.loader_limit)
 )

 self.metadata.steps.append(
 Set(
 fields={
 "recipe_metadata/system_prompt": self.system_prompt,
 "recipe_metadata/format": self.format,
 }
 )
 )

 if self.card:
 self.standardization.steps.extend(self.card.preprocess_steps)

 self.processing.steps.append(self.task)

 if isinstance(self.augmentor, TaskInputsAugmentor):
 self.augmentor.set_fields(self.card.task.augmentable_inputs)
 self.processing.steps.append(self.augmentor)

 if self.has_custom_demos_pool:
 self.processing.steps.append(
 CreateDemosPool(
 from_split=self.demos_taken_from,
 to_split_names=[self.demos_pool_name, self.demos_taken_from],
 to_split_sizes=[int(self.demos_pool_size)],
 remove_targets_from_source_split=self.demos_removed_from_data,
 )
 )

 if self.use_demos:
 if self.sampler is None:
 if self.card.sampler is None:
 raise ValueError(
 "Unexpected None value for card.sampler. "
 "To use num_demos > 0, please set a sampler on the TaskCard."
 )
 self.sampler = self.card.sampler

 self.prepare_refiners()

 if self.use_demos:
 if isinstance(self.num_demos, int):
 self.verbalization.steps.append(
 ConstantSizeSample(
 from_stream=self.demos_pool_name,
 to_field=self.demos_field,
 sampler=self.sampler,
 sample_size=self.num_demos,
 )
 )
 self.verbalization.steps.append(
 Set(fields={"recipe_metadata/num_demos": self.num_demos})
 )

 elif isinstance(self.num_demos, list):
 self.verbalization.steps.append(
 RandomSizeSample(
 from_stream=self.demos_pool_name,
 to_field=self.demos_field,
 sampler=self.sampler,
 sample_sizes=self.num_demos,
 )
 )
 self.verbalization.steps.append(
 GetLength(field="demos", to_field="recipe_metadata/num_demos")
 )
 else:
 raise ValueError("num_demos must be int or List[int]")

 if isinstance(self.template, list):
 self.verbalization.steps.append(
 ApplyRandomTemplate(
 templates=self.template, demos_field=self.demos_field
 )
 )
 else:
 self.verbalization.steps.append(
 ApplySingleTemplate(
 template=self.template, demos_field=self.demos_field
 )
 )
 else:
 self.verbalization.steps.append(
 Set(fields={"recipe_metadata/num_demos": 0})
 )
 if isinstance(self.template, list):
 self.verbalization.steps.append(
 ApplyRandomTemplate(templates=self.template)
 )
 else:
 self.verbalization.steps.append(
 ApplySingleTemplate(template=self.template)
 )

 self.verbalization.steps.append(self.system_prompt)
 self.verbalization.steps.append(self.format)
 if isinstance(self.augmentor, FinalStateInputsAugmentor):
 self.verbalization.steps.append(self.augmentor)

 if self.postprocessors is not None:
 self.finalize.steps.append(
 Set(fields={"postprocessors": self.postprocessors})
 )

 if self.metrics is not None:
 self.finalize.steps.append(Set(fields={"metrics": self.metrics}))

 self.finalize.steps.append(Finalize(group_by=self.group_by))

 def prepare(self):
 if isinstance(self.template, TemplatesList):
 self.template = self.template.items
 self.reset_pipeline()


class StandardRecipeWithIndexes(BaseRecipe):
 template_card_index: int = None

 def prepare(self):
 assert (
 self.template_card_index is None or self.template is None
 ), f"Specify either template ({self.template}) or template_card_index ({self.template_card_index}) but not both"
 assert not (
 self.template_card_index is None and self.template is None
 ), "Specify either template or template_card_index in card"
 if self.template_card_index is not None:
 try:
 self.template = self.card.templates[self.template_card_index]
 except Exception as e:
 if isinstance(self.card.templates, dict):
 options = list(self.card.templates.keys())
 else:
 options = list(range(0, len(self.card.templates)))
 raise ValueError(
 f"card_template_index '{self.template_card_index}' is not defined in card. Possible card_template_index options: {options}"
 ) from e

 super().prepare()


class StandardRecipe(StandardRecipeWithIndexes):
 """This class represents a standard recipe for data processing and preparation.

 This class can be used to prepare a recipe.
 with all necessary steps, refiners and renderers included. It allows to set various
 parameters and steps in a sequential manner for preparing the recipe.

 Attributes:
 card (TaskCard): TaskCard object associated with the recipe.
 template (Template, optional): Template object to be used for the recipe.
 system_prompt (SystemPrompt, optional): SystemPrompt object to be used for the recipe.
 loader_limit (int, optional): Specifies the maximum number of instances per stream to be returned from the loader (used to reduce loading time in large datasets)
 format (SystemFormat, optional): SystemFormat object to be used for the recipe.
 metrics (List[str]): list of catalog metrics to use with this recipe.
 postprocessors (List[str]): list of catalog processors to apply at post processing. (Not recommended to use from here)
 group_by (List[Union[str, List[str]]]): list of task_data or metadata keys to group global scores by.
 train_refiner (StreamRefiner, optional): Train refiner to be used in the recipe.
 max_train_instances (int, optional): Maximum training instances for the refiner.
 validation_refiner (StreamRefiner, optional): Validation refiner to be used in the recipe.
 max_validation_instances (int, optional): Maximum validation instances for the refiner.
 test_refiner (StreamRefiner, optional): Test refiner to be used in the recipe.
 max_test_instances (int, optional): Maximum test instances for the refiner.
 demos_pool_size (int, optional): Size of the demos pool.
 num_demos (int, optional): Number of demos to be used.
 demos_pool_name (str, optional): Name of the demos pool. Default is "demos_pool".
 demos_taken_from (str, optional): Specifies from where the demos are taken. Default is "train".
 demos_field (str, optional): Field name for demos. Default is "demos".
 demos_removed_from_data (bool, optional): whether to remove the demos from the source data, Default is True
 sampler (Sampler, optional): The Sampler used to select the demonstrations when num_demos > 0.
 steps (List[StreamingOperator], optional): List of StreamingOperator objects to be used in the recipe.
 augmentor (Augmentor) : Augmentor to be used to pseudo randomly augment the source text
 instruction_card_index (int, optional): Index of instruction card to be used
 for preparing the recipe.
 template_card_index (int, optional): Index of template card to be used for
 preparing the recipe.

 Methods:
 prepare(): This overridden method is used for preparing the recipe
 by arranging all the steps, refiners, and renderers in a sequential manner.

 Raises:
 AssertionError: If both template and template_card_index are specified at the same time.
 """

 pass