Update _BatchManager to work with GlobalSteps and input_batch_size per step
#366
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
plaguss
reviewed
Feb 29, 2024
| @@ -29,7 +29,7 @@ | |||
| SaveFormats = Literal["json", "yaml"] | |||
|
|
|||
|
|
|||
| def _get_class(module: str = None, name: str = None) -> Type: | |||
| def _get_class(module: str, name: str) -> Type: | |||
There was a problem hiding this comment.
The default args are placed so here we can pass the data using **kwargs. But also, I see I have to update the code to use distilabel/pipeline/serialization and remove this file. Will do that in a different PR.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Description
This PR updates the
_BatchManagerto handle the creation of a "super" batch gathering all the data from the previous batches for theGlobalSteps. In addition, it adds a new attribute to theStepclass calledinput_batch_size. This argument can be used to define the desired number of rows that the step will receive per batch. Having that said, the_BatchManagerhas also been updated to store the data of the output batches, instead of the batches itself. This allows to build batches of the desired size for each step.