Vishwas1 commited on
Commit
67a5c18
·
verified ·
1 Parent(s): 1a39537

Update train_model.py

Browse files
Files changed (1) hide show
  1. train_model.py +15 -3
train_model.py CHANGED
@@ -11,8 +11,17 @@ from datasets import load_dataset, Dataset
11
  import torch
12
  import os
13
  from huggingface_hub import HfApi, HfFolder
 
14
 
15
  def main():
 
 
 
 
 
 
 
 
16
  parser = argparse.ArgumentParser()
17
  parser.add_argument("--task", type=str, required=True, help="Task type: generation or classification")
18
  parser.add_argument("--model_name", type=str, required=True, help="Name of the model")
@@ -24,6 +33,8 @@ def main():
24
  parser.add_argument("--sequence_length", type=int, default=512)
25
  args = parser.parse_args()
26
 
 
 
27
  # Define output directory
28
  output_dir = f"./models/{args.model_name}"
29
  os.makedirs(output_dir, exist_ok=True)
@@ -139,17 +150,18 @@ def main():
139
  tokenizer.save_pretrained(output_dir)
140
 
141
  # Push to Hugging Face Hub
142
- model_repo = f"Vishwas1/{args.model_name}" # Replace 'your-username' with your actual username
143
  try:
144
  api.create_repo(repo_id=model_repo, private=False, token=hf_token)
145
  except Exception as e:
146
- print(f"Repository might already exist: {e}")
147
  model.push_to_hub(model_repo, use_auth_token=hf_token)
148
  tokenizer.push_to_hub(model_repo, use_auth_token=hf_token)
149
 
150
- print(f"Model '{args.model_name}' trained and pushed to Hugging Face Hub at '{model_repo}'.")
151
 
152
  if __name__ == "__main__":
153
  main()
154
 
155
 
 
 
11
  import torch
12
  import os
13
  from huggingface_hub import HfApi, HfFolder
14
+ import logging
15
 
16
  def main():
17
+ # Configure Logging
18
+ logging.basicConfig(
19
+ filename='training.log',
20
+ filemode='a',
21
+ format='%(asctime)s - %(levelname)s - %(message)s',
22
+ level=logging.INFO
23
+ )
24
+
25
  parser = argparse.ArgumentParser()
26
  parser.add_argument("--task", type=str, required=True, help="Task type: generation or classification")
27
  parser.add_argument("--model_name", type=str, required=True, help="Name of the model")
 
33
  parser.add_argument("--sequence_length", type=int, default=512)
34
  args = parser.parse_args()
35
 
36
+ logging.info(f"Starting training for model: {args.model_name}, Task: {args.task}")
37
+
38
  # Define output directory
39
  output_dir = f"./models/{args.model_name}"
40
  os.makedirs(output_dir, exist_ok=True)
 
150
  tokenizer.save_pretrained(output_dir)
151
 
152
  # Push to Hugging Face Hub
153
+ model_repo = f"your-username/{args.model_name}" # Replace 'your-username' with your actual username
154
  try:
155
  api.create_repo(repo_id=model_repo, private=False, token=hf_token)
156
  except Exception as e:
157
+ logging.warning(f"Repository might already exist: {e}")
158
  model.push_to_hub(model_repo, use_auth_token=hf_token)
159
  tokenizer.push_to_hub(model_repo, use_auth_token=hf_token)
160
 
161
+ logging.info(f"Model '{args.model_name}' trained and pushed to Hugging Face Hub at '{model_repo}'.")
162
 
163
  if __name__ == "__main__":
164
  main()
165
 
166
 
167
+