updated

app.py

@@ -2,6 +2,7 @@ from fastapi import FastAPI
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from peft import PeftModel
 import torch
+import os
 
 app = FastAPI()
 
@@ -12,23 +13,31 @@ adapter_path = "thinkingnew/llama_invs_adapter"
 # Check if GPU is available
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load base model
+# Create offload folder if needed
+offload_dir = "./offload"
+os.makedirs(offload_dir, exist_ok=True)
+
+# Load base model with offloading support
 base_model = AutoModelForCausalLM.from_pretrained(
-    base_model_path, torch_dtype=torch.float16 if device == "cuda" else torch.float32, device_map="auto"
-).to(device)
+    base_model_path,
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+    device_map="auto",
+    offload_folder=offload_dir if device == "cpu" else None  # Offload to disk if running on CPU
+)
 
 # Load adapter
-model = PeftModel.from_pretrained(base_model, adapter_path).to(device)
+model = PeftModel.from_pretrained(base_model, adapter_path)
 
 # Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained(base_model_path)
 
-# Load pipeline once (for better performance)
+# Load pipeline once for better performance
 text_pipe = pipeline(
     task="text-generation",
     model=model,
     tokenizer=tokenizer,
-    max_length=512
+    max_length=512,
+    device=0 if device == "cuda" else -1  # Use GPU index 0 if available, otherwise CPU
 )
 
 # Root endpoint for testing