Create modeling.py

utilities/modeling.py

@@ -0,0 +1,86 @@
+from unsloth import FastLanguageModel
+import torch
+from trl import SFTTrainer
+from transformers import TrainingArguments
+from unsloth import is_bfloat16_supported
+
+
+def load_model(model_name, max_seq_length):
+    dtype = None 
+    load_in_4bit = True
+    
+    model, tokenizer = FastLanguageModel.from_pretrained(
+        model_name = model_name,
+        max_seq_length = max_seq_length,
+        dtype = dtype,
+        load_in_4bit = load_in_4bit,
+        # token = ""
+    )
+    return model, tokenizer
+
+
+def get_peft(model, peft, max_seq_length, random_seed):
+    
+    model = FastLanguageModel.get_peft_model(
+    model,
+    r = peft['r',]
+    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
+                      "gate_proj", "up_proj", "down_proj",],
+    lora_alpha = peft['alpha'],
+    lora_dropout = peft['dropout'],
+    bias = peft['bias'], 
+    use_gradient_checkpointing = "unsloth",
+    random_state = random_seed,
+    use_rslora = peft['rslora'],  # We support rank stabilized LoRA
+    loftq_config = peft['loftq_config'], # And LoftQ
+    )
+    return model
+
+
+def get_trainer(model, tokenizer, dataset, sft, data_field, max_seq_length, random_seed):
+    
+    trainer = SFTTrainer(
+        model = model,
+        tokenizer = tokenizer,
+        train_dataset = dataset,
+        dataset_text_field = data_field,
+        max_seq_length = max_seq_length,
+        dataset_num_proc = 2,
+        packing = False,
+        args = TrainingArguments(
+            per_device_train_batch_size = sft['per_device_train_batch_size'],
+            gradient_accumulation_steps = sft['gradient_accumulation_steps'],
+            warmup_steps = sft['warmup_steps'],
+            num_train_epochs = num_epochs,
+            max_steps = max_steps,
+            learning_rate = sft['learning_rate'],
+            fp16 = not is_bfloat16_supported(),
+            bf16 = is_bfloat16_supported(),
+            logging_steps = sft['logging_steps'],
+            optim = sft['optim'],
+            weight_decay = sft['weight_decay'],
+            lr_scheduler_type = sft['lr_scheduler_type'],
+            seed = random_seed,
+            output_dir = "outputs",
+        ),
+    )
+    return trainer
+
+
+def prepare_trainer(model_name, max_seq_length, random_seed,
+                    peft, sft, dataset, data_field):
+
+    print("Loading Model")
+    model, tokenizer = load_model(model_name, max_seq_length)
+
+    print("Preparing for PEFT")
+    model = get_peft(model, peft, max_seq_length, random_seed)
+
+    print("Getting Trainer Model")
+    trainer = get_trainer(model, tokenizer, dataset, data_field, max_seq_length, random_seed)
+
+    return trainer
+
+if __name__ == "__main__":
+    trainer = prepare_trainer()
+