dasds

app.py

@@ -68,12 +68,12 @@ print("Loading model with: quantization_config=", quant_config, ", device_map=",
 # Load model with 8-bit quantization and CPU offloading
 try:
     model = Llama4ForConditionalGeneration.from_pretrained(
-        MODEL_ID,
-        torch_dtype=torch.bfloat16,
-        device_map=device_map,
-        quantization_config=quant_config,
-        attn_implementation="flex_attention"
-    )
+    MODEL_ID,
+    device_map="auto",
+    torch_dtype=torch.float16,
+    quantization_config=quant_config,
+    offload_folder="./offload"
+)
 except Exception as e:
     print(f"Model loading failed: {str(e)}")
     raise

train_llama4.py

@@ -1,7 +1,6 @@
 # train_llama4.py
-# Script to fine-tune Llama 4 Maverick for healthcare fraud detection
 
-from transformers import AutoTokenizer, Llama4ForConditionalGeneration
+from transformers import AutoTokenizer, Llama4ForConditionalGeneration, BitsAndBytesConfig
 import datasets
 import torch
 from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
@@ -9,55 +8,50 @@ from accelerate import Accelerator
 import huggingface_hub
 import os
 
-# Debug: Confirm file version
-print("Running train_llama4.py with CPU offloading (version: 2025-04-21 v2)")
+print("Running train_llama4.py with CPU offloading (version: 2025-04-22 v1)")
 
+# ——————————————————————————
 # Authenticate with Hugging Face
-LLama = os.getenv("LLama")
-if not LLama:
-    raise ValueError("LLama token not found. Set it in Hugging Face Space secrets as 'LLama'.")
-huggingface_hub.login(token=LLama)
+LLAMA = os.getenv("LLama")
+if not LLAMA:
+    raise ValueError("LLama token not found. Set it in environment as 'LLama'.")
+huggingface_hub.login(token=LLAMA)
 
-# Model setup
+# ——————————————————————————
+# Tokenizer
 MODEL_ID = "meta-llama/Llama-4-Maverick-17B-128E-Instruct"
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
-
 if tokenizer.pad_token is None:
     tokenizer.add_special_tokens({'pad_token': '[PAD]'})
 
-# Custom device map for CPU offloading
-device_map = {
-    "model.embed_tokens": 0,
-    "model.layers.0-15": 0,
-    "model.layers.16-31": "cpu",
-    "model.norm": 0,
-    "lm_head": 0
-}
-
-# Debug: Confirm offloading settings
-print("Loading model with CPU offloading: llm_int8_enable_fp32_cpu_offload=True, device_map=", device_map)
+# ——————————————————————————
+# Quantization + CPU off‑load config
+quant_config = BitsAndBytesConfig(
+    load_in_8bit=True,
+    llm_int8_enable_fp32_cpu_offload=True
+)
 
-# Load model with 8-bit quantization and CPU offloading
+print("Loading model with 8-bit quantization, CPU offload, and automatic device mapping")
 model = Llama4ForConditionalGeneration.from_pretrained(
     MODEL_ID,
     torch_dtype=torch.bfloat16,
-    device_map=device_map,
-    quantization_config={"load_in_8bit": True},
-    llm_int8_enable_fp32_cpu_offload=True,
-    attn_implementation="flex_attention"
+    device_map="auto",
+    quantization_config=quant_config,
+    offload_folder="./offload"
 )
 
-# Resize token embeddings
+# Resize embeddings if we added [PAD]
 model.resize_token_embeddings(len(tokenizer))
 
-# Initialize Accelerator
+# ——————————————————————————
+# Prepare for training
 accelerator = Accelerator()
 model = accelerator.prepare(model)
 
-# Load dataset
+# Load training data
 dataset = datasets.load_dataset('json', data_files="Bingaman_training_data.json")['train']
 
-# LoRA configuration
+# LoRA setup
 lora_config = LoraConfig(
     r=16,
     lora_alpha=32,
@@ -67,7 +61,6 @@ lora_config = LoraConfig(
     task_type="CAUSAL_LM"
 )
 
-# Prepare model for fine-tuning
 model = prepare_model_for_kbit_training(model)
 model = get_peft_model(model, lora_config)
 
@@ -87,16 +80,17 @@ training_args = {
     "lr_scheduler_type": "cosine"
 }
 
-# Initialize trainer
+# Initialize Trainer via Accelerate
 trainer = accelerator.prepare(
     datasets.Trainer(
         model=model,
         args=datasets.TrainingArguments(**training_args),
-        train_dataset=dataset,
+        train_dataset=dataset
     )
 )
 
-# Train
+# ——————————————————————————
+# Run training
 trainer.train()
 model.save_pretrained("./fine_tuned_model")
-print("Training completed!")
+print("Training completed!")