Improve generated Notebook

utils/notebook_generator.py

@@ -2,7 +2,7 @@ from typing import Optional
 
 import nbformat as nbf
 
-from utils import FTDataSet, falcon
+from utils import FTDataSet, falcon, gemma
 
 
 def create_install_libraries_cells(cells: list):
@@ -136,6 +136,21 @@ def create_model_cells(cells: list, model_id: str, version: str, flash_attention
         auto_model_import = "FalconForCausalLM"
         trust_code = ""
 
+    chat_ml = """
+# Set chat template to OAI chatML
+model, tokenizer = setup_chat_format(model, tokenizer)
+"""
+    note = f"""
+> **Note:** For `{model_id}`, we will not use `setup_chat_format`. Instead, we will directly use this tokenizer, [philschmid/gemma-tokenizer-chatml](https://huggingface.co/philschmid/gemma-tokenizer-chatml), to fine-tune `{model_id}` with ChatML.
+"""
+    tokenizer_id = f"{model_id}-{version}"
+    if model_id == gemma.name:
+        tokenizer_id = "philschmid/gemma-tokenizer-chatml"
+        chat_ml =""
+    else:
+        note = ""
+
+
     code = f"""
 import torch
 from transformers import AutoTokenizer, {auto_model_import}, BitsAndBytesConfig
@@ -160,15 +175,13 @@ model = {auto_model_import}.from_pretrained(
     quantization_config=bnb_config
 )
 
-tokenizer = AutoTokenizer.from_pretrained(model_id)
+tokenizer = AutoTokenizer.from_pretrained("{tokenizer_id}")
 tokenizer.padding_side = "{pad_side}"
 {pad_value_str}
-
-# Set chat template to OAI chatML
-model, tokenizer = setup_chat_format(model, tokenizer)
+{chat_ml}
     """
 
-    text_1 = """
+    text_1 = f"""
 This process involves two key steps:
 
 1. **LLM Quantization:**
@@ -194,7 +207,9 @@ a 24GB GPU for fine-tuning.
         
         * Sets the tokenizer's chat template, which defines the format used to convert input data into a chat-like 
         structure. The default template is `chatml` from OpenAI.
-        """
+
+{note}
+"""
 
     code_cell = nbf.v4.new_code_cell(code)
     text_cell1 = nbf.v4.new_markdown_cell(text_1)