ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: deduplicate low-risk code paths

Browse Source
This commit is contained in:
ViperEkura
2026-08-19 16:17:40 +08:00
parent f252af495c
commit 398e8a3ea3
12 changed files with 64 additions and 114 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/trainer/strategy.py
+2 -15
View File
@@ -1,6 +1,6 @@
"""Training strategy implementations with factory pattern."""
from abc import ABC, abstractmethod
from abc import ABC
from typing import Callable, Dict, List, Optional, TypedDict, Union
import torch
@@ -187,7 +187,6 @@ class BaseStrategy(ABC):
self.extra_kwargs = kwargs
self._rollout_runner = None
@abstractmethod
def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
"""Compute loss for the given batch.
@@ -197,7 +196,7 @@ class BaseStrategy(ABC):
Returns:
Computed loss tensor
"""
raise NotImplementedError
return self.compute_loss_output(batch)["loss"]
def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput:
return self._normalize_output(self.compute_loss(batch))
@@ -328,9 +327,6 @@ class SEQStrategy(BaseStrategy):
super().__init__(model, device, **kwargs)
self.label_smoothing = label_smoothing
def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
return self.compute_loss_output(batch)["loss"]
def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput:
batch = move_to_device(batch, self.device)
input_ids, target_ids = batch["input_ids"], batch["target_ids"]
@@ -369,9 +365,6 @@ class SFTStrategy(BaseStrategy):
super().__init__(model, device, **kwargs)
self.label_smoothing = label_smoothing
def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
return self.compute_loss_output(batch)["loss"]
def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput:
batch = move_to_device(batch, self.device)
input_ids, target_ids, position_ids, loss_mask = (
@@ -426,9 +419,6 @@ class DPOStrategy(BaseStrategy):
self.beta = beta
self.reduction = reduction
def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
return self.compute_loss_output(batch)["loss"]
def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput:
batch = move_to_device(batch, self.device)
chosen_ids, rejected_ids = batch["chosen"], batch["rejected"]
@@ -553,9 +543,6 @@ class GRPOStrategy(BaseStrategy):
if state_dict is not None:
self.old_model.load_state_dict(state_dict)
def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
return self.compute_loss_output(batch)["loss"]
def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput:
batch = move_to_device(batch, self.device)
prompts = batch["prompts"]