OrcaleSeek / train_advanced.py
prelington's picture
Create train_advanced.py
57a9c7a verified
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",
}