Upload train.py with huggingface_hub

train.py

@@ -11,7 +11,6 @@
 #     "trl",
 #     "peft",
 #     "wandb",
-#     "bitsandbytes",
 #     "torchvision",
 #     "torchaudio",
 # ]
@@ -22,7 +21,7 @@
 
 import torch
 from datasets import load_dataset
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer
 from trl import SFTConfig, SFTTrainer, setup_chat_format
 from peft import LoraConfig
 
@@ -70,20 +69,12 @@ learning_rate = 2e-4
 
 """## Load model and tokenizer"""
 
-# specify how to quantize the model
-# quantization_config = BitsAndBytesConfig(
-#         load_in_4bit=True,
-#         bnb_4bit_quant_type="nf4",
-#         bnb_4bit_use_double_quant=True,
-# )
-
 # Load model
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
     torch_dtype=torch.bfloat16,
     use_cache=False,  # Disable KV cache during training
     device_map="auto",
-    # quantization_config=quantization_config
 )
 
 # Load tokenizer