LMFuser 0.0.1__py3-none-any.whl

lmfuser/schedulers.py

@@ -0,0 +1,210 @@
+from functools import partial
+import math
+
+from torch.optim.lr_scheduler import LambdaLR, LRScheduler
+from torch.optim import Optimizer
+
+from hyperargs import Conf, IntArg, FloatArg, OptionArg, monitor_on, add_dependency
+
+def _get_linear_schedule_with_warmup_lr_lambda(
+    current_step: int, 
+    *, 
+    num_warmup_steps: int, 
+    num_training_steps: int
+):
+    if current_step < num_warmup_steps:
+        return float(current_step) / float(max(1, num_warmup_steps))
+    return max(
+        0.0, 
+        float(num_training_steps - current_step) / float(
+            max(1, num_training_steps - num_warmup_steps))
+        )
+
+def get_linear_schedule_with_warmup(
+    optimizer: Optimizer,
+    num_warmup_steps: int, 
+    num_training_steps: int, 
+    last_epoch=-1
+) -> LRScheduler:
+    """
+    Create a schedule with a learning rate that decreases linearly from the initial lr set in the optimizer to 0, after
+    a warmup period during which it increases linearly from 0 to the initial lr set in the optimizer.
+
+    Args:
+        optimizer ([`~torch.optim.Optimizer`]):
+            The optimizer for which to schedule the learning rate.
+        num_warmup_steps (`int`):
+            The number of steps for the warmup phase.
+        num_training_steps (`int`):
+            The total number of training steps.
+        last_epoch (`int`, *optional*, defaults to -1):
+            The index of the last epoch when resuming training.
+
+    Return:
+        `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
+    """
+
+    lr_lambda = partial(
+        _get_linear_schedule_with_warmup_lr_lambda,
+        num_warmup_steps=num_warmup_steps,
+        num_training_steps=num_training_steps,
+    )
+    return LambdaLR(optimizer, lr_lambda, last_epoch)
+
+def _get_constant_schedule_with_warmup_lr_lambda(
+    current_step: int, 
+    *, 
+    num_warmup_steps: int
+):
+    if current_step < num_warmup_steps:
+        return float(current_step) / float(max(1.0, num_warmup_steps))
+    return 1.0
+
+def get_constant_schedule_with_warmup(
+    optimizer: Optimizer, 
+    num_warmup_steps: int, 
+    last_epoch: int = -1
+) -> LRScheduler:
+    """
+    Create a schedule with a constant learning rate preceded by a warmup period during which the learning rate
+    increases linearly between 0 and the initial lr set in the optimizer.
+
+    Args:
+        optimizer ([`~torch.optim.Optimizer`]):
+            The optimizer for which to schedule the learning rate.
+        num_warmup_steps (`int`):
+            The number of steps for the warmup phase.
+        last_epoch (`int`, *optional*, defaults to -1):
+            The index of the last epoch when resuming training.
+
+    Return:
+        `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
+    """
+
+    lr_lambda = partial(_get_constant_schedule_with_warmup_lr_lambda, num_warmup_steps=num_warmup_steps)
+    return LambdaLR(optimizer, lr_lambda, last_epoch=last_epoch)
+
+def _get_cosine_schedule_with_warmup_lr_lambda(
+    current_step: int, *, num_warmup_steps: int, num_training_steps: int, num_cycles: float
+):
+    if current_step < num_warmup_steps:
+        return float(current_step) / float(max(1, num_warmup_steps))
+    progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
+    return max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress)))
+
+
+def get_cosine_schedule_with_warmup(
+    optimizer: Optimizer, num_warmup_steps: int, num_training_steps: int, num_cycles: float = 0.5, last_epoch: int = -1
+) -> LRScheduler:
+    """
+    Create a schedule with a learning rate that decreases following the values of the cosine function between the
+    initial lr set in the optimizer to 0, after a warmup period during which it increases linearly between 0 and the
+    initial lr set in the optimizer.
+
+    Args:
+        optimizer ([`~torch.optim.Optimizer`]):
+            The optimizer for which to schedule the learning rate.
+        num_warmup_steps (`int`):
+            The number of steps for the warmup phase.
+        num_training_steps (`int`):
+            The total number of training steps.
+        num_cycles (`float`, *optional*, defaults to 0.5):
+            The number of waves in the cosine schedule (the defaults is to just decrease from the max value to 0
+            following a half-cosine).
+        last_epoch (`int`, *optional*, defaults to -1):
+            The index of the last epoch when resuming training.
+
+    Return:
+        `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
+    """
+
+    lr_lambda = partial(
+        _get_cosine_schedule_with_warmup_lr_lambda,
+        num_warmup_steps=num_warmup_steps,
+        num_training_steps=num_training_steps,
+        num_cycles=num_cycles,
+    )
+    return LambdaLR(optimizer, lr_lambda, last_epoch)
+
+
+class LRSchedulerConfigBase(Conf):
+
+    def init_lr_scheduler(self, optimizer: Optimizer, **kwargs) -> LRScheduler:
+        raise NotImplementedError(
+            'Please Implement this method in child classes!'
+        )
+
+
+class ConstantScheduleWithWarmupConfig(LRSchedulerConfigBase):
+
+    num_warmup_steps= IntArg(0, min_value=0)
+    last_epoch= IntArg(-1, min_value=-1)
+
+    def init_lr_scheduler(self, optimizer: Optimizer, **kwargs) -> LRScheduler:
+        return get_constant_schedule_with_warmup(
+            optimizer=optimizer,
+            num_warmup_steps=self.num_warmup_steps.value(), # type: ignore
+            last_epoch=self.last_epoch.value() # type: ignore
+        )
+
+
+class LienarScheduleWithWarmup(LRSchedulerConfigBase):
+
+    num_warmup_steps= IntArg(0, min_value=0)
+    num_training_steps= IntArg(10000, min_value=1)
+    last_epoch= IntArg(-1, min_value=-1)
+
+    def init_lr_scheduler(self, optimizer: Optimizer, **kwargs) -> LRScheduler:
+        return get_linear_schedule_with_warmup(
+            optimizer=optimizer,
+            num_warmup_steps=self.num_warmup_steps.value(), # type: ignore
+            num_training_steps=self.num_training_steps.value(), # type: ignore
+            last_epoch=self.last_epoch.value() # type: ignore
+        )
+
+
+class CosineScheduleWithWarmu(LRSchedulerConfigBase):
+    
+    num_warmup_steps= IntArg(0, min_value=0)
+    num_training_steps= IntArg(10000, min_value=1)
+    num_cycles= FloatArg(0.5, min_value=0.0, max_value=1.0)
+    last_epoch= IntArg(-1, min_value=-1)
+
+    def init_lr_scheduler(self, optimizer: Optimizer, **kwargs) -> LRScheduler:
+        return get_cosine_schedule_with_warmup(
+            optimizer=optimizer,
+            num_warmup_steps=self.num_warmup_steps.value(), # type: ignore
+            num_training_steps=self.num_training_steps.value(), # type: ignore
+            num_cycles=self.num_cycles.value(), # type: ignore
+            last_epoch=self.last_epoch.value() # type: ignore
+        )
+
+
+@add_dependency('type', 'scheduler')
+class LRSchedulerConfig(LRSchedulerConfigBase):
+
+    type = OptionArg(
+        options=[
+            'constant_schedule_with_warmup',
+            'linear_schedule_with_warmup',
+            'cosine_schedule_with_warmup'
+        ],
+        default='constant_schedule_with_warmup',
+    )
+    scheduler = ConstantScheduleWithWarmupConfig()
+
+    def init_lr_scheduler(self, optimizer: Optimizer, **kwargs) -> LRScheduler:
+        return self.scheduler.init_lr_scheduler(optimizer, **kwargs)
+
+    @monitor_on('type')
+    def set_scheduler(self) -> None:
+        if self.type.value() == 'constant_schedule_with_warmup' and not isinstance(self.scheduler, ConstantScheduleWithWarmupConfig):
+            self.scheduler = ConstantScheduleWithWarmupConfig()
+        elif self.type.value() == 'linear_schedule_with_warmup' and not isinstance(self.scheduler, LienarScheduleWithWarmup):
+            self.scheduler = LienarScheduleWithWarmup()
+        elif self.type.value() == 'cosine_schedule_with_warmup' and not isinstance(self.scheduler, CosineScheduleWithWarmu):
+            self.scheduler = CosineScheduleWithWarmu()
+
+if __name__ == '__main__':
+    conf = LRSchedulerConfig.parse_command_line()
+    print(conf)

lmfuser/task.py

@@ -0,0 +1,287 @@
+from abc import ABC
+from typing import TypeVar, Any, Callable
+from collections.abc import Iterable
+
+import torch
+from torch import nn
+from lmfuser_data.interfaces import Batch, Row
+from lmfuser_data.scanners import Scanner
+from lmfuser_data import DataLoader
+from lmfuser_data.interfaces import SubclassTracer
+from hyperargs import Conf, StrArg, FloatArg, IntArg, OptionArg, add_dependency, monitor_on
+
+
+def scanner_type_list() -> list[str]:
+    return list(Scanner.all_subclass_names())
+
+
+@add_dependency('num_train_data_path', 'train_data_path_list')
+@add_dependency('num_train_data_path', 'train_data_weights')
+class TaskBase(Conf, SubclassTracer):
+    num_train_data_path = IntArg(1, min_value=0)
+    train_data_path_list = [StrArg('Enther the path to the data file.')]
+    train_data_weights = [FloatArg(1.0, min_value=0.0, max_value=1.0)]
+
+    num_eval_data_path = IntArg(1, min_value=0)
+    eval_data_path_list = [StrArg('Enther the path to the data file.')]
+    eval_data_weights = [FloatArg(1.0, min_value=0.0, max_value=1.0)]
+
+    scanner_type = OptionArg(default='C4Scanner', option_fn=scanner_type_list)
+
+    _train_dataloader: DataLoader | None = None
+    _eval_dataloader: DataLoader | None = None
+
+    @monitor_on('num_train_data_path')
+    def set_train_path_list(self) -> None:
+        num = self.num_train_data_path.value()
+        assert isinstance(num, int)
+        if len(self.train_data_path_list) > num:
+            self.train_data_path_list = self.train_data_path_list[:num]
+            self.train_data_weights = self.train_data_weights[:num]
+        elif len(self.train_data_path_list) < num:
+            self.train_data_path_list += [StrArg('Enther the path to the data file.')] * (num - len(self.train_data_path_list))
+            self.train_data_weights += [FloatArg(1.0, min_value=0.0, max_value=1.0)] * (num - len(self.train_data_weights))
+
+    @monitor_on('num_eval_data_path')
+    def set_eval_path_list(self) -> None:
+        num = self.num_eval_data_path.value()
+        assert isinstance(num, int)
+        if len(self.eval_data_path_list) > num:
+            self.eval_data_path_list = self.eval_data_path_list[:num]
+            self.eval_data_weights = self.eval_data_weights[:num]
+        elif len(self.eval_data_path_list) < num:
+            self.eval_data_path_list += [StrArg('Enther the path to the data file.')] * (num - len(self.train_data_path_list))
+            self.eval_data_weights += [FloatArg(1.0, min_value=0.0, max_value=1.0)] * (num - len(self.train_data_weights))
+
+    def _get_train_dataloader(
+        self,
+        batch_size: int,
+        seed: int,
+        shuffle: bool,
+        prefetch_factor: int,
+        ignore_error: bool,
+        qps: float | None,
+        instruct_timeout: float,
+        worker_timeout: float,
+        num_workers: int,
+        rank: int,
+        world_size: int
+    ) -> DataLoader | None:
+        if self.num_train_data_path.value() == 0:
+            return None
+        if self._train_dataloader is not None:
+            return self._train_dataloader
+        path_list = [p.value() for p in self.train_data_path_list]
+        weight_list = [w.value() for w in self.train_data_weights]
+        scanner_type = self.scanner_type.value()
+        assert scanner_type is not None, 'scanner_type is None'
+
+        self._train_dataloader = DataLoader(
+            batch_size=batch_size,
+            path_list=path_list, # type: ignore
+            distributor_weights=weight_list, # type: ignore
+            scanner_type=Scanner.get_subclass(scanner_type),
+            seed=seed,
+            shuffle=shuffle,
+            pre_fetch_factor=prefetch_factor,
+            ignore_error=ignore_error,
+            qps=qps,
+            instruct_timeout=instruct_timeout,
+            worker_timeout=worker_timeout,
+            num_workers=num_workers,
+            map_fn=self.get_row_processor(),
+            flow_fn=self.get_flow_processor(),
+            batch_map_fn=self.get_batch_processor(),
+            rank_idx=rank,
+            num_ranks=world_size,
+        )
+
+        return self._train_dataloader
+
+    def _get_eval_dataloader(
+        self,
+        batch_size: int,
+        seed: int,
+        shuffle: bool,
+        prefetch_factor: int,
+        ignore_error: bool,
+        qps: float | None,
+        instruct_timeout: float,
+        worker_timeout: float,
+        num_workers: int,
+        rank: int,
+        world_size: int
+    ) -> DataLoader | None:
+        if self.num_eval_data_path.value() == 0:
+            return None
+        if self._eval_dataloader is not None:
+            return self._eval_dataloader
+        path_list = [p.value() for p in self.eval_data_path_list]
+        weight_list = [w.value() for w in self.eval_data_weights]
+        scanner_type = self.scanner_type.value()
+        assert scanner_type is not None, 'scanner_type is None'
+        self._eval_dataloader = DataLoader(
+            batch_size=batch_size,
+            path_list=path_list, # type: ignore
+            distributor_weights=weight_list, # type: ignore
+            scanner_type=Scanner.get_subclass(scanner_type),
+            seed=seed,
+            shuffle=shuffle,
+            pre_fetch_factor=prefetch_factor,
+            ignore_error=ignore_error,
+            qps=qps,
+            instruct_timeout=instruct_timeout,
+            worker_timeout=worker_timeout,
+            num_workers=num_workers,
+            map_fn=self.get_row_processor(),
+            flow_fn=self.get_flow_processor(),
+            batch_map_fn=self.get_batch_processor(),
+            rank_idx=rank,
+            num_ranks=world_size,
+        )
+
+        return self._eval_dataloader
+
+    def train_step(
+        self, model: nn.Module,
+        batch: Batch,
+        step: int,
+        device: Any,
+        acc_step: int,
+        **kwargs: Any
+    ) -> Batch | torch.Tensor:
+        raise NotImplementedError('Please implement this method in child class')
+
+    def eval_step(
+        self,
+        model: nn.Module,
+        batch: Batch,
+        step: int,
+        device: Any,
+        **kwargs: Any
+    ) -> dict[str, list[Any]]:
+        raise NotImplementedError('Please implement this method in child class')
+
+    def cal_dev_metric(self, eval_outputs: dict[str, list[Any]]) -> dict[str, Any]:
+        raise NotImplementedError('Please implement this method in child class')
+
+    def get_row_processor(self) -> Callable[[Row], Row] | None:
+        return None
+
+    def get_flow_processor(self) -> Callable[[Iterable[Row]], Iterable[Row]] | None:
+        return None
+
+    def get_batch_processor(self) -> Callable[[Batch], Batch] | None:
+        return None
+
+
+class Task(TaskBase):
+    pass
+
+def task_list() -> list[str]:
+    return list(TaskBase.all_subclass_names())
+
+
+@add_dependency('conf', 'task_name')
+class TaskSelector(Conf):
+    task_name = OptionArg(default='Task', option_fn=task_list)
+    conf: TaskBase = Task()
+
+    @monitor_on('task_name')
+    def change_conf(self) -> None:
+        name = self.task_name.value()
+        if name is None:
+            self.conf = Task()
+
+        elif name != self.conf.__class__.__name__:
+            self.conf = TaskBase.all_subclass_map()[name]()
+
+
+@add_dependency('num_tasks', 'tasks')
+class Tasks(Conf):
+    num_tasks = IntArg(1, min_value=1)
+    tasks = [TaskSelector()]
+    task_weights = [FloatArg(1.0, min_value=0.0, max_value=1.0)]
+
+    @monitor_on('num_tasks')
+    def change_task_list(self) -> None:
+        num = self.num_tasks.value()
+        assert num is not None, 'num_tasks is None'
+
+        if len(self.tasks) > num:
+            self.tasks = self.tasks[:num]
+            self.task_weight = self.task_weight[:num]
+        elif len(self.tasks) < num:
+            self.tasks += [TaskSelector() for _ in range(num - len(self.tasks))]
+            self.task_weight += [
+                FloatArg(1.0, min_value=0.0, max_value=1.0) for _ in range(num - len(self.task_weight))
+            ]
+
+    def get_train_dataloaders(
+        self,
+        batch_size: int,
+        seed: int,
+        shuffle: bool,
+        prefetch_factor: int,
+        ignore_error: bool,
+        qps: float | None,
+        instruct_timeout: float,
+        worker_timeout: float,
+        num_workers: int,
+        rank: int,
+        world_size: int
+    ) -> list[DataLoader | None]:
+        return [
+            task.conf._get_train_dataloader(
+                batch_size=batch_size,
+                seed=seed,
+                shuffle=shuffle,
+                prefetch_factor=prefetch_factor,
+                ignore_error=ignore_error,
+                qps=qps,
+                instruct_timeout=instruct_timeout,
+                worker_timeout=worker_timeout,
+                num_workers=num_workers,
+                rank=rank,
+                world_size=world_size,
+            )
+            for task in self.tasks
+        ]
+
+    def get_eval_dataloaders(
+        self,
+        batch_size: int,
+        seed: int,
+        shuffle: bool,
+        prefetch_factor: int,
+        ignore_error: bool,
+        qps: float | None,
+        instruct_timeout: float,
+        worker_timeout: float,
+        num_workers: int,
+        rank: int,
+        world_size: int
+    ) -> list[DataLoader | None]:
+        return [
+            task.conf._get_eval_dataloader(
+                batch_size=batch_size,
+                seed=seed,
+                shuffle=shuffle,
+                prefetch_factor=prefetch_factor,
+                ignore_error=ignore_error,
+                qps=qps,
+                instruct_timeout=instruct_timeout,
+                worker_timeout=worker_timeout,
+                num_workers=num_workers,
+                rank=rank,
+                world_size=world_size,
+            )
+            for task in self.tasks
+        ]
+
+if __name__ == '__main__':
+    class TaskTemp(Task):
+        pass
+
+    conf = Tasks.parse_command_line()
+    print(conf)