File size: 813 Bytes
57a9c7a |
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 |
import logging
import wandb
from transformers import TrainerCallback
# Setup logging
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
handlers=[
logging.FileHandler('training.log'),
logging.StreamHandler()
]
)
class CustomCallback(TrainerCallback):
def on_log(self, args, state, control, logs=None, **kwargs):
if logs:
logging.info(f"Step {state.global_step}: {logs}")
def on_epoch_end(self, args, state, control, **kwargs):
logging.info(f"Epoch {state.epoch} completed")
def setup_wandb():
wandb.init(project="OrcaleSeek", entity="your-username")
wandb.config = {
"learning_rate": 2e-5,
"architecture": "OrcaleSeek",
"dataset": "Your-Dataset",
} |