Update the generated Notebook to login properly to HF

app.py

@@ -2,7 +2,6 @@ from typing import Any
 
 from nbconvert import HTMLExporter
 
-
 from utils.notebook_generator import *
 from utils.components_creator import *
 
@@ -44,6 +43,10 @@ def centered_column():
     return gr.Column(elem_classes=["container"])
 
 
+def should_login_to_hf_model(model_id: str):
+    return model_id == gemma.name or model_id == llama.name
+
+
 def change_model_selection(model_id):
     if model_id == gemma.name:
         gr.Warning("""
@@ -107,10 +110,6 @@ def generate_code(components: dict[Component, Any]):
     if flash_attention_value:
         create_install_flash_attention(notebook['cells'])
 
-    push_to_hub = get_value(components, PUSH_TO_HUB_ID)
-    if push_to_hub:
-        create_login_hf_cells(notebook['cells'])
-
     dataset_value = get_value(components, DATASET_SELECTION_ID)
     seed_value = get_value(components, DATASET_SHUFFLING_SEED)
     if not check_valid_input(dataset_value):
@@ -119,6 +118,8 @@ def generate_code(components: dict[Component, Any]):
         create_datasets_cells(notebook['cells'], get_dataset(dataset_value), seed_value)
 
     model_value = get_value(components, MODEL_SELECTION_ID)
+    should_login = should_login_to_hf_model(model_value)
+
     if not check_valid_input(model_value):
         gr.Warning("No model is selected!")
     else:
@@ -126,6 +127,9 @@ def generate_code(components: dict[Component, Any]):
         if not check_valid_input(version_value):
             gr.Warning("No version of the model is selected")
         else:
+            if should_login:
+                create_login_hf_cells(notebook['cells'], should_login=True, model_name=model_value)
+
             load_in_4bit = get_value(components, LOAD_IN_4_BIT_ID)
             bnb_4bit_use_double_quant = get_value(components, BNB_4BIT_USE_DOUBLE_QUANT)
             bnb_4bit_quant_type = get_value(components, BNB_4BIT_QUANT_TYPE)
@@ -174,6 +178,8 @@ def generate_code(components: dict[Component, Any]):
     packing = get_value(components, PACKING_ID)
     create_sft_trainer_cells(notebook['cells'], max_seq_length, packing)
 
+    push_to_hub = get_value(components, PUSH_TO_HUB_ID)
+
     create_start_training_cells(notebook['cells'], epochs, max_steps, push_to_hub, output_dir)
 
     create_free_gpu_cells(notebook['cells'])
@@ -181,6 +187,8 @@ def generate_code(components: dict[Component, Any]):
     create_merge_lora_cells(notebook['cells'], output_dir)
 
     if push_to_hub:
+        if not should_login:
+            create_login_hf_cells(notebook['cells'])
         push_merged_model_cells(notebook['cells'], output_dir)
 
     file_name = f"{finetuning_notebook}.ipynb"

utils/notebook_generator.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 import nbformat as nbf
 
 from utils import FTDataSet
@@ -47,7 +49,7 @@ def create_install_flash_attention(cells: list):
 import torch; assert torch.cuda.get_device_capability()[0] >= 8, 'Hardware not supported for Flash Attention'
 
 !pip install ninja packaging
-!MAX_JOBS=4 pip install flash-attn --no-build-isolation --upgrade
+!MAX_JOBS=4 pip install -q flash-attn --no-build-isolation --upgrade
     """
     code_cell = nbf.v4.new_code_cell(code)
     cells.append(text_cell)
@@ -55,13 +57,16 @@ import torch; assert torch.cuda.get_device_capability()[0] >= 8, 'Hardware not s
     cells.append(code_cell)
 
 
-def create_login_hf_cells(cells: list):
+def create_login_hf_cells(cells: list, should_login: bool = False, model_name: Optional[str] = None):
     text_cell = nbf.v4.new_markdown_cell(
         "### Login to HF")
-    text_cell1 = nbf.v4.new_markdown_cell("Installing **huggingface_hub** to use as a remote "
-                                          "model versioning service. This means that your model, logs, and information "
-                                          "will be automatically pushed to the Hub during training. You should have "
-                                          "'HF_TOKEN'")
+
+    text_1 = "Login with your `HF_TOKEN` in order to push the finetuned model to `huggingface_hub`."
+
+    if should_login:
+        text_1 = f"Login with your `HF_TOKEN` in order to load **{model_name}** from `huggingface_hub`."
+
+    text_cell1 = nbf.v4.new_markdown_cell(text_1)
     code = """
 # Install huggingface_hub
 !pip install -q huggingface_hub
@@ -229,6 +234,8 @@ def create_training_args_cells(cells: list, epochs, max_steps, logging_steps, pe
     elif report_to != "none":
         to_install = report_to
 
+    gradient_checkpointing_kwargs = {"use_reentrant": False}
+
     code_report = f"""
 # Installing {to_install} to report the metrics   
     
@@ -244,12 +251,12 @@ args = TrainingArguments(
     per_device_train_batch_size={per_device_train_batch_size},
     gradient_accumulation_steps={gradient_accumulation_steps},
     gradient_checkpointing={gradient_checkpointing},
+    gradient_checkpointing_kwargs={gradient_checkpointing_kwargs},
     optim="adamw_torch_fused",
     logging_steps={logging_steps},
     save_strategy='{save_strategy}',
     learning_rate={learning_rate},
     bf16=True,
-    tf32=True,
     max_grad_norm={max_grad_norm},
     warmup_ratio={warmup_ratio},
     lr_scheduler_type='{lr_scheduler_type}',