Update app.py

app.py

@@ -1,5 +1,5 @@
-# updated_app.py
-# Enhanced Gradio app for Llama 4 Maverick healthcare fraud detection (text-only)
+# app.py
+# Enhanced Gradio app for Llama 4 Maverick healthcare fraud detection (text-only with CPU offloading)
 
 import gradio as gr
 from transformers import AutoTokenizer, Llama4ForConditionalGeneration
@@ -45,12 +45,22 @@ tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
 if tokenizer.pad_token is None:
     tokenizer.add_special_tokens({'pad_token': '[PAD]'})
 
-# Load model with 8-bit quantization to fit in 80 GB VRAM
+# Custom device map to offload some layers to CPU
+device_map = {
+    "model.embed_tokens": 0,
+    "model.layers.0-15": 0,  # Keep first 16 layers on GPU
+    "model.layers.16-31": "cpu",  # Offload remaining layers to CPU
+    "model.norm": 0,
+    "lm_head": 0
+}
+
+# Load model with 8-bit quantization and CPU offloading
 model = Llama4ForConditionalGeneration.from_pretrained(
     MODEL_ID,
     torch_dtype=torch.bfloat16,
-    device_map="auto",
+    device_map=device_map,
     quantization_config={"load_in_8bit": True},
+    llm_int8_enable_fp32_cpu_offload=True,
     attn_implementation="flex_attention"
 )