Update app.py

app.py

@@ -34,29 +34,29 @@ with open('config.json', 'r') as config_file:
 num_args = len(config)
 
 
-arg2 = config.get('arg2', '1') 
-arg3 = config.get('arg3', 'train_fleet') 
-arg4 = config.get('arg4', 'test_fleet')
-arg5 = config.get('arg5', '1') 
-arg6 = config.get('arg6', '1')
-arg7 = config.get('arg7', 'saved_fleet_model')
-
-if num_args == 7:    
+arg1 = config.get('arg1', '1') 
+arg2 = config.get('arg2', 'train_fleet') 
+arg3 = config.get('arg3', 'test_fleet')
+arg4 = config.get('arg4', '1') 
+arg5 = config.get('arg5', '1')
+arg6 = config.get('arg6', 'saved_fleet_model')
+
+if num_args == 6:    
     # cmd args
     # sys.argv[0] is the script name, sys.argv[1] is the first argument, etc.
-    should_train_model = arg2  # should train model?
-    train_file = arg3   # training file name
-    test_file = arg4    # eval file name
-    batch_size_for_trainer = int(arg5)  # batch sizes to send to trainer
-    should_produce_eval_matrix = int(arg6)     # should produce matrix?
-    path_to_save_trained_model_to = arg7
-
-    print(f"should train model? : {arg2}")
-    print (f"file to train on : {arg3}")
-    print (f"file to evaluate on : {arg4}")
-    print (f"batch size : {arg5}")
-    print (f"should produce eval matrix : {arg6}")
-    print (f"path to save trained model : {arg7}")
+    should_train_model = arg1  # should train model?
+    train_file = arg2   # training file name
+    test_file = arg3    # eval file name
+    batch_size_for_trainer = int(arg4)  # batch sizes to send to trainer
+    should_produce_eval_matrix = int(arg5)     # should produce matrix?
+    path_to_save_trained_model_to = arg6
+
+    print(f"should train model? : {arg1}")
+    print (f"file to train on : {arg2}")
+    print (f"file to evaluate on : {arg3}")
+    print (f"batch size : {arg4}")
+    print (f"should produce eval matrix : {arg5}")
+    print (f"path to save trained model : {arg6}")
     
     print(f"should train model? : {should_train_model}")
     print (f"file to train on : {train_file}")
@@ -331,27 +331,27 @@ if (should_train_model=='1'): #train model
     # Upload the model and tokenizer to the Hugging Face repository
     
     upload_folder(
-        folder_path=f"{model_save_path}",
-        path_in_repo=f"{model_save_path}",
+        folder_path=f"{model_save_path}_model",
+        path_in_repo=f"{model_save_path}_model",
         repo_id=repo_name,
         token=api_token,
-        commit_message="Push fleet model",
+        commit_message="Push model",
         #overwrite=True  # Force overwrite existing files
     )
 
     upload_folder(
-        folder_path="saved_fleet_tokenizer",
-        path_in_repo="saved_fleet_tokenizer",
+        folder_path=f"{model_save_path}_tokenizer",
+        path_in_repo=f"{model_save_path}_tokenizer",
         repo_id=repo_name,
         token=api_token,
-        commit_message="Push fleet tokenizer",
+        commit_message="Push tokenizer",
         #overwrite=True  # Force overwrite existing files
     )
     
 else:
     print('Load Pre-trained')
-    model_save_path = "./saved_fleet_model"
-    tokenizer_save_path = "./saved_fleet_tokenizer"
+    model_save_path = f"./{model_save_path}_model"
+    tokenizer_save_path = f"./{model_save_path}_tokenizer"
     # RobertaTokenizer.from_pretrained(model_save_path)
     model = AutoModelForSequenceClassification.from_pretrained(model_save_path).to('cpu')
     tokenizer = AutoTokenizer.from_pretrained(tokenizer_save_path)