Update train_llama4.py

train_llama4.py

@@ -1,7 +1,7 @@
 # train_llama4.py
-# Script to fine-tune Llama 4 Maverick for healthcare fraud detection
+# Script to fine-tune Llama 4 Maverick for healthcare fraud detection (text-only)
 
-from transformers import AutoProcessor, Llama4ForConditionalGeneration, Trainer, TrainingArguments
+from transformers import AutoTokenizer, Llama4ForConditionalGeneration, Trainer, TrainingArguments
 from transformers import BitsAndBytesConfig
 import datasets
 import torch
@@ -22,9 +22,13 @@ if not LLama:
     raise ValueError("LLama token not found. Set it in Hugging Face Space secrets as 'LLama'.")
 huggingface_hub.login(token=LLama)
 
-# Load Llama 4 model and processor
+# Load Llama 4 model and tokenizer
 MODEL_ID = "meta-llama/Llama-4-Maverick-17B-128E-Instruct"
-processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+
+# Add padding token if it doesn't exist
+if tokenizer.pad_token is None:
+    tokenizer.add_special_tokens({'pad_token': '[PAD]'})
 
 # Quantization config for A100 80 GB VRAM
 quantization_config = BitsAndBytesConfig(load_in_8bit=True)
@@ -56,18 +60,8 @@ print("First example from dataset:", dataset["train"][0])
 
 # Tokenization
 def tokenize_data(example):
-    messages = [
-        {
-            "role": "user",
-            "content": [{"type": "text", "text": example['input']}]
-        },
-        {
-            "role": "assistant",
-            "content": [{"type": "text", "text": example['output']}]
-        }
-    ]
-    formatted_text = processor.apply_chat_template(messages, add_generation_prompt=False)
-    inputs = processor(formatted_text, padding="max_length", truncation=True, max_length=4096, return_tensors="pt")
+    formatted_text = f"<s>[INST] {example['input']} [/INST] {example['output']}</s>"
+    inputs = tokenizer(formatted_text, padding="max_length", truncation=True, max_length=4096, return_tensors="pt")
     input_ids = inputs["input_ids"].squeeze(0).tolist()
     attention_mask = inputs["attention_mask"].squeeze(0).tolist()
     labels = input_ids.copy()
@@ -124,5 +118,5 @@ trainer = Trainer(
 # Start training
 trainer.train()
 model.save_pretrained("./fine_tuned_llama4_healthcare")
-processor.save_pretrained("./fine_tuned_llama4_healthcare")
-print("Training complete. Model and processor saved to ./fine_tuned_llama4_healthcare")
+tokenizer.save_pretrained("./fine_tuned_llama4_healthcare")
+print("Training complete. Model and tokenizer saved to ./fine_tuned_llama4_healthcare")