Update finetune_flan_t5.py

finetune_flan_t5.py

@@ -1,35 +1,34 @@
 from datasets import load_dataset
 from transformers import (
-    AutoTokenizer, 
-    AutoModelForSeq2SeqLM, 
+    T5ForConditionalGeneration,  # Using specific model class
+    AutoTokenizer,
     TrainingArguments,
-    DataCollatorForSeq2Seq,
-    FlaxAutoModelForSeq2SeqLM  # Added for explicit model loading
+    DataCollatorForSeq2Seq
 )
 from trl import SFTTrainer
 import torch
 
-# 1. Load and prepare dataset
+
+
+# 2. Load and prepare dataset
 dataset = load_dataset("json", data_files="data/med_q_n_a_converted.jsonl", split="train")
 
-# Add formatted text field
-dataset = dataset.map(lambda x: {
-    "text": f"### Instruction:\n{x['input']}\n\n### Response:\n{x['output']}"
-})
+# Create properly formatted text field
+def format_example(example):
+    return {
+        "text": f"Instruction: {example['input']}\nResponse: {example['output']}",
+        "input": example["input"],
+        "output": example["output"]
+    }
+
+dataset = dataset.map(format_example)
 
-# 2. Load model and tokenizer - METHOD 1: Explicit FLAN-T5 loading
+# 3. Load model and tokenizer
 model_name = "google/flan-t5-base"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-
-# METHOD 1: Load model directly without AutoModel
-from transformers import T5ForConditionalGeneration
 model = T5ForConditionalGeneration.from_pretrained(model_name)
 
-# METHOD 2: Or install Japanese support (if needed)
-# pip install transformers[ja]
-# Then use AutoModel as before
-
-# 3. Training arguments
+# 4. Configure training
 training_args = TrainingArguments(
     output_dir="./flan-t5-medical-finetuned",
     per_device_train_batch_size=4,
@@ -41,22 +40,26 @@ training_args = TrainingArguments(
     evaluation_strategy="no",
     fp16=torch.cuda.is_available(),
     report_to="none",
-    remove_unused_columns=False
+    remove_unused_columns=False,
+    # Add these to prevent version conflicts
+    dataloader_pin_memory=False,
+    dataloader_num_workers=0
 )
 
-# 4. Initialize trainer
+# 5. Initialize trainer with proper config
 trainer = SFTTrainer(
     model=model,
     tokenizer=tokenizer,
     train_dataset=dataset,
     args=training_args,
     dataset_text_field="text",
+    max_seq_length=512,  # Explicitly set to avoid warning
     data_collator=DataCollatorForSeq2Seq(
         tokenizer,
         model=model,
-        padding=True
+        padding="longest"
     )
 )
 
-# 5. Start training
+# 6. Start training
 trainer.train()