Update finetune_flan_t5.py

finetune_flan_t5.py

@@ -1,75 +1,54 @@
 from datasets import load_dataset
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, TrainingArguments
-from trl import SFTTrainer  # ✅ from trl
-from transformers import DataCollatorForSeq2Seq  # ✅ from transformers
+from transformers import (
+    AutoTokenizer, 
+    AutoModelForSeq2SeqLM, 
+    TrainingArguments,
+    DataCollatorForSeq2Seq
+)
+from trl import SFTTrainer
 import torch
 
-# Load your dataset (from the converted JSONL file)
+# 1. Load dataset
 dataset = load_dataset("json", data_files="data/med_q_n_a_converted.jsonl", split="train")
 
-# Load tokenizer and model
+# 2. Load model and tokenizer
 model_name = "google/flan-t5-base"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
 
-# Preprocess dataset
-def preprocess(example):
-    input_text = example["input"]
-    target_text = example["output"]
-
-    model_inputs = tokenizer(
-        input_text,
-        max_length=512,
-        truncation=True,
-        padding="max_length"
-    )
-
-    labels = tokenizer(
-        target_text,
-        max_length=128,
-        truncation=True,
-        padding="max_length"
-    )["input_ids"]
-
-    model_inputs["labels"] = labels
-    return model_inputs
-
-
+# 3. Formatting function for SFTTrainer
+def format_instruction(example):
+    return f"### Instruction:\n{example['input']}\n\n### Response:\n{example['output']}"
 
-# Apply preprocessing
-tokenized_dataset = dataset.map(preprocess)
-
-
-# Define training arguments
+# 4. Training arguments
 training_args = TrainingArguments(
-    output_dir="./flan-t5-medical",
+    output_dir="./flan-t5-medical-finetuned",
     per_device_train_batch_size=4,
     gradient_accumulation_steps=2,
     num_train_epochs=3,
+    learning_rate=5e-5,
     logging_dir="./logs",
     save_strategy="epoch",
     evaluation_strategy="no",
-    fp16=torch.cuda.is_available()
+    fp16=torch.cuda.is_available(),
+    report_to="none"
 )
 
-# Define data collator
-data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)
-
-def formatting_func(example):
-    return [f"Input: {example['input']}\nOutput: {example['output']}"]
-
-from trl import SFTTrainer
-from transformers import DataCollatorForSeq2Seq
-
+# 5. Initialize SFTTrainer correctly
 trainer = SFTTrainer(
     model=model,
     tokenizer=tokenizer,
-    train_dataset=tokenized_dataset,  # already tokenized
+    train_dataset=dataset,
     args=training_args,
-    data_collator=DataCollatorForSeq2Seq(tokenizer, padding=True),
-    packing=False,
-    tokenized_dataset=True  # ✅ Now supported after upgrade
+    max_seq_length=512,
+    formatting_func=format_instruction,
+    data_collator=DataCollatorForSeq2Seq(
+        tokenizer,
+        pad_to_multiple_of=8,
+        return_tensors="pt",
+        padding=True
+    )
 )
 
-# Start training
-trainer.train()
+# 6. Start training
+trainer.train()