Update app.py

app.py

@@ -1,23 +1,26 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
 from fastapi.responses import JSONResponse
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from accelerate import init_empty_weights, infer_auto_device_map, dispatch_model
 
 print("Version ---- 1")
 app = FastAPI()
 
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
-
 # Load model and tokenizer from Hugging Face
 model_name = "deepseek-ai/DeepSeek-R1-Distill-Llama-8B"
 
 tokenizer = AutoTokenizer.from_pretrained(model_name)
+
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
-    torch_dtype=torch.bfloat16,  # Use bfloat16 if supported, else change to torch.float16
-    device_map="auto"            # Automatically use GPU if available
+    dtype=torch.bfloat16,
+    device_map="auto",
+    offload_folder="offload"   # folder for disk offload
 )
 
+
 class ConflictDetectionRequest(BaseModel):
     Req1: str
     Req2: str