Update app.py

app.py

@@ -8,6 +8,9 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, LogitsProcessor, LogitsProcessorList
 from peft import PeftModel
 
+OFFLOAD_DIR = "offload"
+os.makedirs(OFFLOAD_DIR, exist_ok=True)
+
 
 # ----------------------------
 # 1. CONFIG
@@ -63,15 +66,24 @@ if tokenizer.pad_token_id is None:
 
 def load_lora(repo_id: str):
     """Load one LoRA adapter on top of the base model."""
+    # Base model, let HF/accelerate place layers (GPU + CPU)
     base = AutoModelForCausalLM.from_pretrained(
         BASE_MODEL,
         torch_dtype=DTYPE,
         device_map="auto",
     )
-    model = PeftModel.from_pretrained(base, repo_id)
+
+    # IMPORTANT: give accelerate an offload folder
+    model = PeftModel.from_pretrained(
+        base,
+        repo_id,
+        device_map="auto",
+        offload_folder=OFFLOAD_DIR,   # <-- fixes the offload_dir error
+    )
     model.eval()
     return model
 
+
 print("Loading LoRA models (this happens once at startup)...")
 model_confusion   = load_lora(HF_CONFUSION)
 model_engagement  = load_lora(HF_ENGAGEMENT)
@@ -198,8 +210,11 @@ def make_prompt(label_name: str, means: dict):
 
 def run_one_model(model, label_name: str, means: dict) -> int:
     prompt = make_prompt(label_name, means)
+    # figure out the right device (for auto-sharded models this is e.g. "cuda:0")
+    device = getattr(model, "device", DEVICE)
+
     with torch.no_grad():
-        toks = tokenizer(prompt, return_tensors="pt").to(DEVICE)
+        toks = tokenizer(prompt, return_tensors="pt").to(device)
         out = model.generate(
             **toks,
             max_new_tokens=1,
@@ -208,11 +223,15 @@ def run_one_model(model, label_name: str, means: dict) -> int:
             pad_token_id=tokenizer.pad_token_id,
             eos_token_id=tokenizer.eos_token_id,
         )
-        text = tokenizer.decode(out[0, toks["input_ids"].shape[1]:], skip_special_tokens=True)
+        text = tokenizer.decode(
+            out[0, toks["input_ids"].shape[1]:],
+            skip_special_tokens=True,
+        )
     m = re.search(r"[0-3]", text)
     return int(m.group()) if m else -1
 
 
+
 # ----------------------------
 # 5. GRADIO PIPELINE
 # ----------------------------