Update app.py

app.py

@@ -1,20 +1,27 @@
 from fastapi import FastAPI, Query
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
+import os
 
 app = FastAPI()
 
-# Load the model and tokenizer
+# Create offload folder if not exists
+os.makedirs("./offload", exist_ok=True)
+
+# Load tokenizer and model with offload_folder to prevent device_map error
 tokenizer = AutoTokenizer.from_pretrained("deepseek-ai/deepseek-llm-7b-base")
 model = AutoModelForCausalLM.from_pretrained(
     "deepseek-ai/deepseek-llm-7b-base",
     torch_dtype=torch.float16,
-    device_map="auto"
+    device_map="auto",
+    offload_folder="./offload"
 )
 
 @app.get("/")
 def home():
-    return {"message": "DeepSeek LLM is running. Use /ask?prompt=your+question"}
+    return {
+        "message": "✅ DeepSeek LLM is running. Use endpoint /ask?prompt=your+question"
+    }
 
 @app.get("/ask")
 def ask(prompt: str = Query(..., description="Your input prompt")):