update requirements.txt

Author: HaomingX

baselines/src/finetune.py

@@ -1,7 +1,6 @@
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig, set_seed, Trainer
 
-import wandb
 import transformers
 import os
 from peft import LoraConfig, get_peft_model
@@ -10,8 +9,6 @@
 from src.utils import get_model_identifiers_from_yaml, find_all_linear_names
 from src.dataset import QADataset, DefaultDataset
 
-os.environ['WANDB_MODE'] = 'dryrun'
-
 def finetune(cfg):
     if os.environ.get('LOCAL_RANK') is not None:
         local_rank = int(os.environ.get('LOCAL_RANK', '0'))
@@ -28,12 +25,6 @@ def finetune(cfg):
     model_cfg = get_model_identifiers_from_yaml(cfg.model_family)
     model_id = model_cfg["hf_key"]
 
-    wandb.init(project='finetune', config={
-    "learning_rate": cfg.lr,
-    "epochs": cfg.num_epochs,
-    "batch_size": batch_size * gradient_accumulation_steps * num_devices,
-}, name=f'finetune-lr{cfg.lr}-epoch{cfg.num_epochs}')
-
     Path(cfg.save_dir).mkdir(parents=True, exist_ok=True)
     # save the cfg file
     #if master process

requirements.txt

@@ -16,4 +16,5 @@ scipy==1.14.1
 ninja==1.11.1.2
 zhipuai==2.1.5.20241203
 openai==1.55.3
-vllm==0.6.5
+vllm==0.6.5
+scikit-learn==1.6.1