Spaces:
Sleeping
Sleeping
File size: 1,817 Bytes
8a35bc0 7c4332a 264e02e 7c4332a 264e02e 7c4332a 264e02e 7c4332a 264e02e 7c4332a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 |
# -------------------------------------------------------------------
# Pimcore
#
# This source file is available under two different licenses:
# - GNU General Public License version 3 (GPLv3)
# - Pimcore Commercial License (PCL)
# Full copyright and license information is available in
# LICENSE.md which is distributed with this source code.
#
# @copyright Copyright (c) Pimcore GmbH (http://www.pimcore.org)
# @license http://www.pimcore.org/license GPLv3 and PCL
# -------------------------------------------------------------------
import logging
from transformers import TrainerCallback, TrainingArguments, TrainerState, TrainerControl
from .training_status import TrainingStatus
logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)
class ProgressCallback(TrainerCallback):
__trainingStatus: TrainingStatus = None
__startPercentage: int = None
__endPercentage: int = None
def __init__(self, trainingStatus: TrainingStatus, startPercentage: int, endPercentage: int):
self.__trainingStatus = trainingStatus
self.__startPercentage = startPercentage
self.__endPercentage = endPercentage
def on_step_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
logger.info(f"Completed step {state.global_step} of {state.max_steps}")
if self.__trainingStatus.is_training_aborted():
control.should_training_stop = True
logger.info("Training aborted")
return
scope = self.__endPercentage - self.__startPercentage
progress = round(self.__startPercentage + (state.global_step / state.max_steps) * scope, 2)
self.__trainingStatus.update_status(progress, f"Training model, completed step {state.global_step} of {state.max_steps}")
|