Create automatically a README.md (HF model card)

app.py

@@ -2,6 +2,7 @@ from typing import Any
 
 from nbconvert import HTMLExporter
 
+from utils.create_info_files import create_hf_card
 from utils.notebook_generator import *
 from utils.components_creator import *
 
@@ -68,10 +69,6 @@ def change_model_selection(model_id):
     return None
 
 
-def handle_push_to_hub(value):
-    return gr.Textbox(visible=value)
-
-
 def check_valid_input(value):
     if isinstance(value, str):
         return value and value.strip()
@@ -124,6 +121,7 @@ def generate_code(components: dict[Component, Any]):
     model_value = get_value(components, MODEL_SELECTION_ID)
     should_login = should_login_to_hf_model(model_value)
 
+    version_value = ""
     if not check_valid_input(model_value):
         gr.Warning("No model is selected!")
     else:
@@ -192,6 +190,12 @@ def generate_code(components: dict[Component, Any]):
 
     merge_model_cells(notebook['cells'], output_dir)
 
+    create_readme = get_value(components, README_ID)
+    if create_readme:
+        create_hf_card(notebook['cells'], name=output_dir, base_model_name=model_value,
+                       base_model_version=version_value,
+                       dataset_name=dataset_value, output_dir=output_dir, report_to=report_to)
+
     if push_to_hub:
         if not should_login:
             create_login_hf_cells(notebook['cells'])
@@ -287,8 +291,6 @@ with gr.Blocks(css=css, theme=gr.themes.Soft(text_size='lg', font=["monospace"],
         with centered_column():
             output_dir_cmp, push_to_hub_cmp = add_outputs()
             all_components.update({output_dir_cmp, push_to_hub_cmp})
-            repo_name_cmp = add_hf_repo_cmp()
-            all_components.update({repo_name_cmp})
         with centered_column():
             all_components.update(add_outputs1())
 
@@ -317,12 +319,6 @@ with gr.Blocks(css=css, theme=gr.themes.Soft(text_size='lg', font=["monospace"],
         outputs=version_selection
     )
 
-    push_to_hub_cmp.change(
-        fn=handle_push_to_hub,
-        inputs=push_to_hub_cmp,
-        outputs=repo_name_cmp
-    )
-
 demo.launch(allowed_paths=["/"])
 
 # Upload metrics to the hub....

utils/__init__.py

@@ -41,6 +41,8 @@ REPOSITORY_NAME_ID = "repo_id"
 
 REPORT_TO_ID = "report_to"
 
+README_ID = "readme"
+
 MAX_SEQ_LENGTH_ID = "max_seq_length"
 PACKING_ID = "packing"
 

utils/components_creator.py

@@ -201,8 +201,13 @@ def add_outputs1() -> Set[Component]:
              "'comet_ml', 'mlflow', 'tensorboard' and 'wandb'. Use 'all' to report to all integrations installed, "
              "'none' for no integrations."
     )
+    create_readme = gr.Checkbox(label="Automatically Generate a README.md", value=True, interactive=True,
+                                info="Choose whether to automatically generate a model card (README.md) or not.",
+                                elem_id=README_ID)
+
     out_components: Set[Component] = set()
     out_components.add(report_to)
+    out_components.add(create_readme)
     return out_components
 
 

utils/create_info_files.py

@@ -0,0 +1,74 @@
+import nbformat as nbf
+
+
+def create_hf_card(cells, name, base_model_name, base_model_version, dataset_name, output_dir, report_to):
+    text = f"""
+card = '''
+---
+license: apache-2.0
+tags:
+- generated_from_trainer
+- {base_model_name}
+- PyTorch
+- transformers
+- trl
+- peft
+- {report_to}
+base_model: {base_model_name}-{base_model_version}
+widget:
+  - example_title: Pirate!
+    messages:
+      - role: system
+        content: You are a pirate chatbot who always responds with Arr!
+      - role: user
+        content: "There's a llama on my lawn, how can I get rid of him?"
+    output:
+      text: >-
+        Arr! 'Tis a puzzlin' matter, me hearty! A llama on yer lawn be a rare
+        sight, but I've got a plan that might help ye get rid of 'im. Ye'll need
+        to gather some carrots and hay, and then lure the llama away with the
+        promise of a tasty treat. Once he's gone, ye can clean up yer lawn and
+        enjoy the peace and quiet once again. But beware, me hearty, for there
+        may be more llamas where that one came from! Arr!
+model-index:
+- name: {name}
+  results: []
+datasets:
+- {dataset_name}
+language:
+- en
+pipeline_tag: text-generation
+---
+
+# Model Card for {name}:
+
+**{name}** is a language model that is trained to act as helpful assistant. It is a finetuned version of [{base_model_name}-{base_model_version}](https://huggingface.co/{base_model_name}-{base_model_version}) that was trained using SFTTrainer on of publicly available dataset [
+{dataset_name}](https://huggingface.co/datasets/{dataset_name}).
+
+## Training Procedure:
+
+The training code used to create this model was generated by [Menouar/LLM-FineTuning-Notebook-Generator](https://huggingface.co/spaces/Menouar/LLM-FineTuning-Notebook-Generator).
+
+
+
+## Training hyperparameters
+
+The following hyperparameters were used during the training:
+
+
+'''
+
+with open("{output_dir}/README.md", "w") as f:
+    f.write(card)
+
+args_dict = vars(args)
+
+with open("{output_dir}/README.md", "a") as f:
+    for k, v in args_dict.items():
+        f.write(f"- {{k}}: {{v}}")
+        f.write("\\n \\n")
+"""
+    title = """### Generating a model card (README.md)"""
+    cells.append(nbf.v4.new_markdown_cell(title))
+    code_cell = nbf.v4.new_code_cell(text)
+    cells.append(code_cell)

utils/notebook_generator.py

@@ -61,10 +61,10 @@ def create_login_hf_cells(cells: list, should_login: bool = False, model_name: O
     text_cell = nbf.v4.new_markdown_cell(
         "### Login to HF")
 
-    text_1 = "Login with your `HF_TOKEN` in order to push the finetuned model to `huggingface_hub`."
+    text_1 = "Login with our `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_1 = f"Login with our `HF_TOKEN` in order to load **{model_name}** from `huggingface_hub`."
 
     text_cell1 = nbf.v4.new_markdown_cell(text_1)
     code = """
@@ -167,7 +167,7 @@ This process involves two key steps:
 
 1. **LLM Quantization:**
     - We first load the selected large language model (LLM).
-    - We then use the "bitsandbytes" library to quantize the model, which can significantly reduce its memory footprint.
+    - We then use the `bitsandbytes` library to quantize the model, which can significantly reduce its memory footprint.
 
 > **Note:** The memory requirements of the model scale with its size. For instance, a 7B parameter model may require 
 a 24GB GPU for fine-tuning. 
@@ -228,7 +228,7 @@ def create_training_args_cells(cells: list, epochs, max_steps, logging_steps, pe
                                save_strategy, gradient_accumulation_steps, gradient_checkpointing,
                                learning_rate, max_grad_norm, warmup_ratio, lr_scheduler_type, output_dir,
                                report_to, seed):
-    text_cell = nbf.v4.new_markdown_cell("### TrainingArguments")
+    text_cell = nbf.v4.new_markdown_cell("### Setting the TrainingArguments")
     to_install = None
     if report_to == "all":
         to_install = "azure_ml comet_ml mlflow tensorboard wandb"
@@ -238,8 +238,7 @@ def create_training_args_cells(cells: list, epochs, max_steps, logging_steps, pe
     gradient_checkpointing_kwargs = {"use_reentrant": False}
 
     code_report = f"""
-# Installing {to_install} to report the metrics   
-    
+# Installing {to_install} to report the metrics
 !pip install -q {to_install}
     """
 
@@ -278,7 +277,7 @@ args = TrainingArguments(
 
 def create_sft_trainer_cells(cells: list, max_seq_length, packing):
     text_cell = nbf.v4.new_markdown_cell(
-        """### Supervised Finetuning Trainer (SFT Trainer)
+        """### Setting the Supervised Finetuning Trainer (`SFTTrainer`)
     
 This `SFTTrainer` is a wrapper around the `transformers.Trainer` class and inherits all of its attributes and methods.
 The trainer takes care of properly initializing the `PeftModel`.   
@@ -308,7 +307,7 @@ trainer = SFTTrainer(
 
 def create_start_training_cells(cells: list, epochs, max_steps, push_to_hub, output_dir):
     if push_to_hub:
-        save_txt = "and to the hub."
+        save_txt = f"and to the hub in 'User/{output_dir}'."
     else:
         save_txt = "."
 
@@ -320,7 +319,7 @@ def create_start_training_cells(cells: list, epochs, max_steps, push_to_hub, out
         f"""### Starting Training and Saving Model/Tokenizer
 
 We start training the model by calling the `train()` method on the trainer instance. This will start the training 
-loop and train the model for `{epoch_str}`. The model will be automatically saved the output directory(temp_{output_dir})
+loop and train the model for `{epoch_str}`. The model will be automatically saved to the output directory ('temp_{output_dir}')
 {save_txt} 
   
     """)
@@ -342,7 +341,8 @@ trainer.save_model()
 
 def create_free_gpu_cells(cells: list):
     text_cell = nbf.v4.new_markdown_cell(
-        """### Free the GPU Memory for Merging `PeftModel`""")
+        """### Free the GPU Memory to Prepare for the Merging of the `PerfModel`
+""")
 
     code = f"""
 
@@ -358,11 +358,11 @@ torch.cuda.empty_cache()
 
 def create_merge_lora_cells(cells: list, output_dir):
     text_cell = nbf.v4.new_markdown_cell(
-        """### Merge LoRA adapter in to the Original Model
+        """### Merging LoRa Adapters into the Original Model
 
 While utilizing `LoRA`, we focus on training the adapters rather than the entire model. Consequently, during the 
 model saving process, only the `adapter weights` are preserved, not the complete model. If we wish to save the 
-entire model for easier usage with Text Generation Inference, you can incorporate the adapter weights into the model 
+entire model for easier usage with Text Generation Inference, we can incorporate the adapter weights into the model 
 weights. This can be achieved using the `merge_and_unload` method. Following this, the model can be saved using the 
 `save_pretrained` method. The result is a default model that is ready for inference.
 """)
@@ -378,7 +378,7 @@ model = AutoPeftModelForCausalLM.from_pretrained(
     low_cpu_mem_usage=True
 )
     
-# Merge LoRA and base model and save
+# Merge LoRA with the base model and save
 merged_model = model.merge_and_unload()
 merged_model.save_pretrained("{output_dir}", safe_serialization=True, max_shard_size="2GB")
 tokenizer.save_pretrained("{output_dir}")
@@ -396,23 +396,17 @@ def merge_model_cells(cells: list, output_dir):
 import os
 import shutil
 
-# Specify the source folder and the destination folder
 source_folder = "temp_{output_dir}"
 destination_folder = "{output_dir}"
 
-# Create the destination folder if it doesn't exist
 os.makedirs(destination_folder, exist_ok=True)
 
-# Iterate over the files and subfolders in the source folder
 for item in os.listdir(source_folder):
     item_path = os.path.join(source_folder, item)
     
-    # Check if it's a subfolder (and not a file)
     if os.path.isdir(item_path):
-        # Specify the destination path
         destination_path = os.path.join(destination_folder, item)
         
-        # Copy the subfolder to the destination folder
         shutil.copytree(item_path, destination_path)
 """
 
@@ -422,7 +416,7 @@ for item in os.listdir(source_folder):
 
 
 def push_to_hub_cells(cells: list, output_dir):
-    text = f"Push '{output_dir}' to your Hugging Face account."
+    text = f"### Pushing '{output_dir}' to our Hugging Face account."
     code = f"""
 from huggingface_hub import HfApi, HfFolder, Repository
 
@@ -437,8 +431,7 @@ repo = api.create_repo(token=HfFolder.get_token(), repo_type="model", repo_id=re
 
 api.upload_folder(
     folder_path="{output_dir}",
-    repo_id=repo.repo_id,
-    repo_type="model"
+    repo_id=repo.repo_id
 )
 """
     code_cell = nbf.v4.new_code_cell(code)