Merge branch 'main' of https://huggingface.co/spaces/Cylanoid/Nursing-Home-Fraud-Detection-using-Llama

requirements.txt

@@ -7,4 +7,5 @@ peft==0.14.0
 bitsandbytes
 sentencepiece
 huggingface_hub>=0.19
-accelerate
+accelerate
+flash-attn

train_llama.py

@@ -22,12 +22,13 @@ if tokenizer.pad_token is None:
 # Quantization config
 quantization_config = BitsAndBytesConfig(load_in_8bit=True)
 
-# Load model without FlashAttention
+# Load model with FlashAttention 2
 model = LlamaForCausalLM.from_pretrained(
     MODEL_ID,
     torch_dtype=torch.bfloat16,
     device_map="auto",
-    quantization_config=quantization_config
+    quantization_config=quantization_config,
+    attn_implementation="flash_attention_2"
 )
 
 # Prepare for LoRA
@@ -43,16 +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 fixed length
+# Tokenization with lists (no tensors)
 def tokenize_data(example):
     formatted_text = f"{example['input']} {example['output']}"
-    inputs = tokenizer(formatted_text, truncation=True, max_length=512, padding="max_length", return_tensors="pt")
-    input_ids = inputs["input_ids"].squeeze(0).tolist()
-    attention_mask = inputs["attention_mask"].squeeze(0).tolist()
+    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] * input_len
-    print(f"Debug: input_ids[:5] = {input_ids[:5]}, labels[:5] = {labels[:5]}, attention_mask[:5] = {attention_mask[:5]}")
     return {
         "input_ids": input_ids,
         "labels": labels,
@@ -60,18 +60,24 @@ def tokenize_data(example):
     }
 
 tokenized_dataset = dataset["train"].map(tokenize_data, batched=False, remove_columns=dataset["train"].column_names)
+# Print first example (lists with lengths)
 first_example = tokenized_dataset[0]
 print("First tokenized example:", {k: (type(v), len(v)) for k, v in first_example.items()})
 
-# Data collator with tensor stacking
+# 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(input_ids),
-        "attention_mask": torch.stack(attention_mask),
-        "labels": torch.stack(labels)
+        "input_ids": input_ids,
+        "attention_mask": attention_mask,
+        "labels": labels
     }
 
 # Accelerator and training