oh boy oh noooo

train_llama.py

@@ -2,6 +2,7 @@ from transformers import LlamaForCausalLM, LlamaTokenizer, Trainer, TrainingArgu
 from transformers import BitsAndBytesConfig
 import datasets
 import torch
+from torch.nn.utils.rnn import pad_sequence
 from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
 from accelerate import Accelerator
 
@@ -43,15 +44,15 @@ model.print_trainable_parameters()
 dataset = datasets.load_dataset("json", data_files="final_combined_fraud_data.json", field="training_pairs")
 print("First example from dataset:", dataset["train"][0])
 
-# Tokenization with tensors
+# Tokenization with lists (no tensors)
 def tokenize_data(example):
     formatted_text = f"{example['input']} {example['output']}"
-    inputs = tokenizer(formatted_text, truncation=True, max_length=2048, return_tensors="pt")
-    input_ids = inputs["input_ids"].squeeze(0)
-    labels = inputs["input_ids"].clone().squeeze(0)
+    inputs = tokenizer(formatted_text, truncation=True, max_length=2048)
+    input_ids = inputs["input_ids"]
+    attention_mask = inputs["attention_mask"]
+    labels = input_ids.copy()
     input_len = len(tokenizer(example['input'])["input_ids"])
-    labels[:input_len] = -100
-    attention_mask = inputs["attention_mask"].squeeze(0)
+    labels[:input_len] = [-100] * input_len
     return {
         "input_ids": input_ids,
         "labels": labels,
@@ -59,16 +60,24 @@ def tokenize_data(example):
     }
 
 tokenized_dataset = dataset["train"].map(tokenize_data, batched=False, remove_columns=dataset["train"].column_names)
-# Fix print to handle potential list or tensor
+# Print first example (lists with lengths)
 first_example = tokenized_dataset[0]
-print("First tokenized example:", {k: (type(v), v.shape if hasattr(v, 'shape') else len(v)) for k, v in first_example.items()})
+print("First tokenized example:", {k: (type(v), len(v)) for k, v in first_example.items()})
 
-# Data collator
+# Data collator: convert lists to tensors and pad
 def custom_data_collator(features):
+    input_ids = [torch.tensor(f["input_ids"]) for f in features]
+    attention_mask = [torch.tensor(f["attention_mask"]) for f in features]
+    labels = [torch.tensor(f["labels"]) for f in features]
+    
+    input_ids = pad_sequence(input_ids, batch_first=True, padding_value=tokenizer.pad_token_id)
+    attention_mask = pad_sequence(attention_mask, batch_first=True, padding_value=0)
+    labels = pad_sequence(labels, batch_first=True, padding_value=-100)
+    
     return {
-        "input_ids": torch.stack([f["input_ids"] for f in features]),
-        "attention_mask": torch.stack([f["attention_mask"] for f in features]),
-        "labels": torch.stack([f["labels"] for f in features])
+        "input_ids": input_ids,
+        "attention_mask": attention_mask,
+        "labels": labels
     }
 
 # Accelerator and training