Tonic commited on
Commit
dfcb060
·
1 Parent(s): fcf2981

fixes custom trackio implementation

Browse files
Files changed (1) hide show
  1. scripts/training/train_gpt_oss.py +10 -5
scripts/training/train_gpt_oss.py CHANGED
@@ -12,7 +12,6 @@ import torch
12
  from transformers import AutoTokenizer, AutoModelForCausalLM
13
  from peft import LoraConfig, get_peft_model
14
  from trl import SFTTrainer, SFTConfig
15
- import trackio
16
  from datasets import load_dataset
17
 
18
  def load_gpt_oss_model_and_tokenizer(config):
@@ -85,10 +84,16 @@ def setup_trackio_tracking(config):
85
 
86
  print(f"Setting up Trackio tracking: {config.trackio_url}")
87
 
88
- # Initialize Trackio client
89
- trackio_client = trackio.Client(
90
- api_url=config.trackio_url,
91
- token=config.trackio_token
 
 
 
 
 
 
92
  )
93
 
94
  return trackio_client
 
12
  from transformers import AutoTokenizer, AutoModelForCausalLM
13
  from peft import LoraConfig, get_peft_model
14
  from trl import SFTTrainer, SFTConfig
 
15
  from datasets import load_dataset
16
 
17
  def load_gpt_oss_model_and_tokenizer(config):
 
84
 
85
  print(f"Setting up Trackio tracking: {config.trackio_url}")
86
 
87
+ # Import the correct TrackioAPIClient
88
+ import sys
89
+ import os
90
+ sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'trackio_tonic'))
91
+ from trackio_api_client import TrackioAPIClient
92
+
93
+ # Initialize Trackio client using the correct API
94
+ trackio_client = TrackioAPIClient(
95
+ space_id=config.trackio_url,
96
+ hf_token=config.trackio_token
97
  )
98
 
99
  return trackio_client