updated adfe

train_llama.py

@@ -24,7 +24,7 @@ quantization_config = BitsAndBytesConfig(load_in_8bit=True)
 # Load model with FlashAttention 2
 model = LlamaForCausalLM.from_pretrained(
     MODEL_ID,
-    torch_dtype=torch.bfloat16,  # Matches A100
+    torch_dtype=torch.bfloat16,
     device_map="auto",
     quantization_config=quantization_config,
     attn_implementation="flash_attention_2"
@@ -43,7 +43,7 @@ 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 (dynamic padding)
+# Tokenization with tensors
 def tokenize_data(example):
     formatted_text = f"{example['input']} {example['output']}"
     inputs = tokenizer(formatted_text, truncation=True, max_length=2048, return_tensors="pt")
@@ -51,10 +51,15 @@ def tokenize_data(example):
     labels = inputs["input_ids"].clone().squeeze(0)
     input_len = len(tokenizer(example['input'])["input_ids"])
     labels[:input_len] = -100
-    return {"input_ids": input_ids, "labels": labels, "attention_mask": inputs["attention_mask"].squeeze(0)}
+    attention_mask = inputs["attention_mask"].squeeze(0)
+    return {
+        "input_ids": input_ids,
+        "labels": labels,
+        "attention_mask": attention_mask
+    }
 
 tokenized_dataset = dataset["train"].map(tokenize_data, batched=False, remove_columns=dataset["train"].column_names)
-print("First tokenized example:", {k: (type(v), v.shape if isinstance(v, torch.Tensor) else "list") for k, v in tokenized_dataset[0].items()})
+print("First tokenized example:", {k: (type(v), v.shape) for k, v in tokenized_dataset[0].items()})
 
 # Data collator
 def custom_data_collator(features):