we w

app.py

@@ -2,7 +2,7 @@
 # Gradio app for Llama 4 Maverick healthcare fraud detection (text-only with CPU offloading)
 
 import gradio as gr
-from transformers import AutoTokenizer, Llama4ForConditionalGeneration
+from transformers import AutoTokenizer, Llama4ForConditionalGeneration, BitsAndBytesConfig
 import datasets
 import torch
 import json
@@ -24,7 +24,7 @@ except LookupError:
 from document_analyzer import HealthcareFraudAnalyzer
 
 # Debug: Confirm file version
-print("Running updated app.py with CPU offloading (version: 2025-04-21 v2)")
+print("Running updated app.py with CPU offloading (version: 2025-04-21 v3)")
 
 # Debug: Print environment variables
 print("Environment variables:", dict(os.environ))
@@ -47,27 +47,36 @@ 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
+# Explicit quantization configuration
+quant_config = BitsAndBytesConfig(
+    load_in_8bit=True,
+    llm_int8_enable_fp32_cpu_offload=True
+)
+
+# Custom device map for CPU offloading (more layers to CPU)
 device_map = {
     "model.embed_tokens": 0,
-    "model.layers.0-15": 0,  # First 16 layers on GPU
-    "model.layers.16-31": "cpu",  # Remaining layers on CPU
+    "model.layers.0-10": 0,  # First 11 layers on GPU
+    "model.layers.11-31": "cpu",  # Remaining layers on 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)
+print("Loading model with: quantization_config=", quant_config, ", device_map=", device_map)
 
 # Load model with 8-bit quantization and CPU offloading
-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"
-)
+try:
+    model = Llama4ForConditionalGeneration.from_pretrained(
+        MODEL_ID,
+        torch_dtype=torch.bfloat16,
+        device_map=device_map,
+        quantization_config=quant_config,
+        attn_implementation="flex_attention"
+    )
+except Exception as e:
+    print(f"Model loading failed: {str(e)}")
+    raise
 
 # Resize token embeddings if pad token was added
 model.resize_token_embeddings(len(tokenizer))