Use FalconForCausalLM for falcon finetuning instead of AutoModelForCausalLM

utils/notebook_generator.py

@@ -2,7 +2,7 @@ from typing import Optional
 
 import nbformat as nbf
 
-from utils import FTDataSet
+from utils import FTDataSet, falcon
 
 
 def create_install_libraries_cells(cells: list):
@@ -130,9 +130,13 @@ def create_model_cells(cells: list, model_id: str, version: str, flash_attention
     if pad_value is None:
         pad_value_str = ""
 
+    auto_model_import = "AutoModelForCausalLM"
+    if model_id == falcon.name:
+        auto_model_import = "FalconForCausalLM"
+
     code = f"""
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+from transformers import AutoTokenizer, {auto_model_import}, BitsAndBytesConfig
 from trl import setup_chat_format
 
 # Hugging Face model id
@@ -145,7 +149,7 @@ bnb_config = BitsAndBytesConfig(
 )
 
 # Load model and tokenizer
-model = AutoModelForCausalLM.from_pretrained(
+model = {auto_model_import}.from_pretrained(
     model_id,
     device_map="auto",
     trust_remote_code=True,