Update finetune_flan_t5.py

finetune_flan_t5.py

@@ -8,20 +8,22 @@ from transformers import (
 from trl import SFTTrainer
 import torch
 
-# 1. Load dataset
+# 1. Load and prepare dataset
 dataset = load_dataset("json", data_files="data/med_q_n_a_converted.jsonl", split="train")
 
+# Add 'text' field containing the formatted examples
+def add_text_field(example):
+    example['text'] = f"### Instruction:\n{example['input']}\n\n### Response:\n{example['output']}"
+    return example
+
+dataset = dataset.map(add_text_field)
+
 # 2. Load model and tokenizer
 model_name = "google/flan-t5-base"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
 
-# 3. CORRECTED Formatting function - returns single string per example
-def format_instruction(example):
-    # Return a single formatted string
-    return f"### Instruction:\n{example['input']}\n\n### Response:\n{example['output']}"
-
-# 4. Training arguments
+# 3. Training arguments
 training_args = TrainingArguments(
     output_dir="./flan-t5-medical-finetuned",
     per_device_train_batch_size=4,
@@ -35,23 +37,22 @@ training_args = TrainingArguments(
     report_to="none"
 )
 
-# 5. Initialize SFTTrainer with correct parameters
+# 4. Initialize SFTTrainer with correct configuration
 trainer = SFTTrainer(
     model=model,
     tokenizer=tokenizer,
     train_dataset=dataset,
     args=training_args,
     max_seq_length=512,
-    formatting_func=format_instruction,  # Now returns single string
+    dataset_text_field="text",  # Field we created
     data_collator=DataCollatorForSeq2Seq(
         tokenizer,
-        model=model,  # Added model reference
+        model=model,
         pad_to_multiple_of=8,
         return_tensors="pt",
         padding=True
-    ),
-    dataset_text_field="text"  # Explicit field name
+    )
 )
 
-# 6. Start training
+# 5. Start training
 trainer.train()