Spaces:
Sleeping
Sleeping
Update train_model.py
Browse files- 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"
|
143 |
try:
|
144 |
api.create_repo(repo_id=model_repo, private=False, token=hf_token)
|
145 |
except Exception as e:
|
146 |
-
|
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 |
-
|
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 |
+
|