geekyrakshit commited on
Commit
780c9f0
·
1 Parent(s): 2900eb1

add: wandb job_type

Browse files
guardrails_genie/train/llama_guard.py CHANGED
@@ -213,6 +213,7 @@ class LlamaGuardFineTuner:
213
  project=self.wandb_project,
214
  entity=self.wandb_entity,
215
  name=f"{self.model_name}-{self.dataset_name}",
 
216
  )
217
  self.model.classifier = nn.Linear(
218
  self.model.classifier.in_features, num_classes
 
213
  project=self.wandb_project,
214
  entity=self.wandb_entity,
215
  name=f"{self.model_name}-{self.dataset_name}",
216
+ job_type="fine-tune-llama-guard",
217
  )
218
  self.model.classifier = nn.Linear(
219
  self.model.classifier.in_features, num_classes
guardrails_genie/train/train_classifier.py CHANGED
@@ -62,7 +62,12 @@ def train_binary_classifier(
62
  Exception: If an error occurs during training, the exception is raised after
63
  ensuring Weights & Biases run is finished.
64
  """
65
- wandb.init(project=project_name, entity=entity_name, name=run_name)
 
 
 
 
 
66
  if streamlit_mode:
67
  st.markdown(
68
  f"Explore your training logs on [Weights & Biases]({wandb.run.url})"
 
62
  Exception: If an error occurs during training, the exception is raised after
63
  ensuring Weights & Biases run is finished.
64
  """
65
+ wandb.init(
66
+ project=project_name,
67
+ entity=entity_name,
68
+ name=run_name,
69
+ job_type="train-binary-classifier",
70
+ )
71
  if streamlit_mode:
72
  st.markdown(
73
  f"Explore your training logs on [Weights & Biases]({wandb.run.url})"