Update app.py

app.py

@@ -102,6 +102,7 @@ def test(base_model_id, dataset):
     ###################
     # Hyper-parameters
     ###################
+    print("111")
     training_config = {
         "bf16": True,
         "do_eval": False,
@@ -125,7 +126,8 @@ def test(base_model_id, dataset):
         "gradient_accumulation_steps": 1,
         "warmup_ratio": 0.2,
         }
-    
+
+    print("222")
     peft_config = {
         "r": 16,
         "lora_alpha": 32,
@@ -142,6 +144,7 @@ def test(base_model_id, dataset):
     ###############
     # Setup logging
     ###############
+    print("333")
     logging.basicConfig(
         format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
         datefmt="%Y-%m-%d %H:%M:%S",
@@ -156,6 +159,7 @@ def test(base_model_id, dataset):
     transformers.utils.logging.enable_explicit_format()
     
     # Log on each process a small summary
+    print("444")
     logger.warning(
         f"Process rank: {train_conf.local_rank}, device: {train_conf.device}, n_gpu: {train_conf.n_gpu}"
         + f" distributed training: {bool(train_conf.local_rank != -1)}, 16-bits training: {train_conf.fp16}"
@@ -167,6 +171,7 @@ def test(base_model_id, dataset):
     ################
     # Model Loading
     ################
+    print("444")
     checkpoint_path = "microsoft/Phi-3-mini-4k-instruct"
     # checkpoint_path = "microsoft/Phi-3-mini-128k-instruct"
     model_kwargs = dict(
@@ -176,6 +181,7 @@ def test(base_model_id, dataset):
         torch_dtype=torch.bfloat16,
         device_map=None
     )
+    print("555")
     model = AutoModelForCausalLM.from_pretrained(checkpoint_path, **model_kwargs)
     tokenizer = AutoTokenizer.from_pretrained(checkpoint_path)
     tokenizer.model_max_length = 2048
@@ -187,10 +193,8 @@ def test(base_model_id, dataset):
     ##################
     # Data Processing
     ##################
-    def apply_chat_template(
-        example,
-        tokenizer,
-    ):
+    print("666")
+    def apply_chat_template(example, tokenizer):
         messages = example["messages"]
         example["text"] = tokenizer.apply_chat_template(
             messages, tokenize=False, add_generation_prompt=False)
@@ -200,7 +204,8 @@ def test(base_model_id, dataset):
     train_dataset = raw_dataset["train_sft"]
     test_dataset = raw_dataset["test_sft"]
     column_names = list(train_dataset.features)
-    
+
+    print("777")
     processed_train_dataset = train_dataset.map(
         apply_chat_template,
         fn_kwargs={"tokenizer": tokenizer},
@@ -208,7 +213,8 @@ def test(base_model_id, dataset):
         remove_columns=column_names,
         desc="Applying chat template to train_sft",
     )
-    
+
+    print("888")
     processed_test_dataset = test_dataset.map(
         apply_chat_template,
         fn_kwargs={"tokenizer": tokenizer},
@@ -221,6 +227,7 @@ def test(base_model_id, dataset):
     ###########
     # Training
     ###########
+    print("999")
     trainer = SFTTrainer(
         model=model,
         args=train_conf,
@@ -242,6 +249,7 @@ def test(base_model_id, dataset):
     #############
     # Evaluation
     #############
+    print("aaa")
     tokenizer.padding_side = 'left'
     metrics = trainer.evaluate()
     metrics["eval_samples"] = len(processed_test_dataset)
@@ -252,6 +260,7 @@ def test(base_model_id, dataset):
     # ############
     # # Save model
     # ############
+    print("bbb")
     trainer.save_model(train_conf.output_dir)
 
 def download_model(base_model_id):