Update train_model.py

train_model.py

@@ -11,8 +11,17 @@ from datasets import load_dataset, Dataset
 import torch
 import os
 from huggingface_hub import HfApi, HfFolder
+import logging
 
 def main():
+    # Configure Logging
+    logging.basicConfig(
+        filename='training.log',
+        filemode='a',
+        format='%(asctime)s - %(levelname)s - %(message)s',
+        level=logging.INFO
+    )
+
     parser = argparse.ArgumentParser()
     parser.add_argument("--task", type=str, required=True, help="Task type: generation or classification")
     parser.add_argument("--model_name", type=str, required=True, help="Name of the model")
@@ -24,6 +33,8 @@ def main():
     parser.add_argument("--sequence_length", type=int, default=512)
     args = parser.parse_args()
     
+    logging.info(f"Starting training for model: {args.model_name}, Task: {args.task}")
+
     # Define output directory
     output_dir = f"./models/{args.model_name}"
     os.makedirs(output_dir, exist_ok=True)
@@ -139,17 +150,18 @@ def main():
     tokenizer.save_pretrained(output_dir)
     
     # Push to Hugging Face Hub
-    model_repo = f"Vishwas1/{args.model_name}"  # Replace 'your-username' with your actual username
+    model_repo = f"your-username/{args.model_name}"  # Replace 'your-username' with your actual username
     try:
         api.create_repo(repo_id=model_repo, private=False, token=hf_token)
     except Exception as e:
-        print(f"Repository might already exist: {e}")
+        logging.warning(f"Repository might already exist: {e}")
     model.push_to_hub(model_repo, use_auth_token=hf_token)
     tokenizer.push_to_hub(model_repo, use_auth_token=hf_token)
     
-    print(f"Model '{args.model_name}' trained and pushed to Hugging Face Hub at '{model_repo}'.")
+    logging.info(f"Model '{args.model_name}' trained and pushed to Hugging Face Hub at '{model_repo}'.")
     
 if __name__ == "__main__":
     main()
 
 
+