1 files changed, 13 insertions, 9 deletions

diff --git a/finetune.py b/finetune.py
index bca7ac1..05b558f 100644
--- a/finetune.py
+++ b/finetune.py
@@ -1,15 +1,17 @@
 import torch
 from datasets import load_dataset
-from peft import LoraConfig
-from transformers import AutoModelForCausalLM, AutoModelForSeq2SeqLM, AutoTokenizer
+from peft import LoraConfig, get_peft_model
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 from trl import SFTTrainer, SFTConfig
 import os
 
-#MODEL_ID = "google/gemma-3-270m-it"
-MODEL_ID = "google/gemma-3-1b-it"
+#MODEL_ID = "google/flan-t5-large"
+#MODEL_ID = "google/t5gemma-b-b-prefixlm-it"
+#MODEL_ID = "google/long-t5-tglobal-base"
+MODEL_ID = "google/t5gemma-s-s-ul2-it"
 DATASET_TRAIN_PATH = "train.jsonl"
 DATASET_VAL_PATH = "validation.jsonl"
-OUTPUT_DIR = "./gemma3-1b-it-sponsors"
+OUTPUT_DIR = "./t5gemma-s-s-ul2-it-sponsors"
 LORA_R = 16
 LORA_ALPHA = 32
 LORA_DROPOUT = 0.05
@@ -24,7 +26,7 @@ def main():
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
 
-    model = AutoModelForCausalLM.from_pretrained(
+    model = AutoModelForSeq2SeqLM.from_pretrained(
         MODEL_ID,
         device_map="auto",
         torch_dtype=torch.bfloat16,
@@ -44,9 +46,13 @@ def main():
         target_modules=LORA_TARGET_MODULES,
         lora_dropout=LORA_DROPOUT,
         bias="none",
-        task_type="CAUSAL_LM"
+        task_type="SEQ_2_SEQ_LM"
     )
 
+    print("Applying LoRA config to the model...")
+    model = get_peft_model(model, peft_config)
+    model.print_trainable_parameters()
+
     print("Model and LoRA config loaded. Setting up trainer...")
 
     training_args = SFTConfig(
@@ -57,7 +63,6 @@ def main():
         per_device_eval_batch_size=1,
         per_device_train_batch_size=1,
         gradient_accumulation_steps=8,
-        gradient_checkpointing=True,
         learning_rate=1e-4,
         lr_scheduler_type="linear",
         optim="adamw_torch",
@@ -83,7 +88,6 @@ def main():
         model=model,
         processing_class=tokenizer,
         args=training_args,
-        peft_config=peft_config,
         train_dataset=train_dataset,
         eval_dataset=validation_dataset,
     )