Update app.py

app.py

@@ -16,21 +16,25 @@ SYSTEM_PROMPT = (
     "No use English unless person ask am."
 )
 
+# Hugging Face Spaces-safe writable dir for disk offload
+OFFLOAD_DIR = os.getenv("OFFLOAD_DIR", "/tmp/offload")
+
 
 def load_model():
-    # Tokenizer (use base tokenizer)
+    # Tokenizer (base)
     tokenizer = AutoTokenizer.from_pretrained(
         BASE_MODEL,
         trust_remote_code=True,
     )
 
     if not torch.cuda.is_available():
-        raise RuntimeError(
-            "CUDA GPU not detected. gpt-oss-20b needs a GPU for this demo."
-        )
+        raise RuntimeError("CUDA GPU not detected. gpt-oss-20b needs a GPU for this demo.")
+
+    # Make sure offload folder exists (required when device_map triggers disk offload)
+    os.makedirs(OFFLOAD_DIR, exist_ok=True)
 
-    # If the base model is MXFP4-quantized, DO NOT use BitsAndBytes.
-    # Dequantize=True makes this run on non-H100 GPUs too (L4/A10/T4 etc).
+    # MXFP4 model: do NOT use BitsAndBytes.
+    # dequantize=True allows running on non-H100 GPUs too (L4/A10/T4 etc).
     qconfig = Mxfp4Config(dequantize=True)
 
     base = AutoModelForCausalLM.from_pretrained(
@@ -39,6 +43,8 @@ def load_model():
         torch_dtype=torch.bfloat16,
         quantization_config=qconfig,
         trust_remote_code=True,
+        offload_folder=OFFLOAD_DIR,
+        offload_state_dict=True,
     )
 
     model = PeftModel.from_pretrained(base, ADAPTER_ID)
@@ -49,10 +55,8 @@ def load_model():
 tokenizer, model = load_model()
 
 
-def build_prompt(message, history):
-    # Keep prompt bounded: last N turns
-    MAX_TURNS = 8
-    history = history[-MAX_TURNS:]
+def build_prompt(message, history, max_turns=8):
+    history = (history or [])[-max_turns:]
 
     lines = [SYSTEM_PROMPT, ""]
     for u, a in history:
@@ -84,8 +88,6 @@ def chat(message, history, max_new_tokens, temperature, top_p):
     )
 
     decoded = tokenizer.decode(out[0], skip_special_tokens=True)
-
-    # Extract only the latest assistant segment
     reply = decoded.split("Assistant:")[-1].strip()
     return reply
 
@@ -101,4 +103,4 @@ demo = gr.ChatInterface(
     description=f"Base: {BASE_MODEL} | Adapter: {ADAPTER_ID}",
 )
 
-demo.launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", "7860")))
+demo.launch(server_name="0.0.0.0", server_port=int(os.getenv('PORT', '7860')))