Update app.py

app.py

@@ -4,8 +4,9 @@ from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import os
 
-# Create a folder for offloading weights if needed
-os.makedirs("offload", exist_ok=True)
+# Use a writable folder for offloading weights (Hugging Face Spaces restricts /app)
+offload_dir = "/tmp/offload"
+os.makedirs(offload_dir, exist_ok=True)
 
 app = FastAPI()
 
@@ -18,19 +19,19 @@ app.add_middleware(
     allow_headers=["*"]
 )
 
-# Model name
+# Model name (7B model - large, will offload to /tmp)
 model_name = "ethzanalytics/RedPajama-INCITE-7B-Base-sharded-bf16"
 
 # Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-# Load model with offload folder
+# Load model with /tmp offload folder
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
     torch_dtype="bfloat16",
     device_map="auto",
     low_cpu_mem_usage=True,
-    offload_folder="offload"  # Required after HF update
+    offload_folder=offload_dir
 )
 
 # Request body schema