Update app.py

app.py

@@ -1,6 +1,6 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 import uvicorn
 
@@ -9,19 +9,15 @@ app = FastAPI()
 # Model name (update with your actual model path on Hugging Face)
 model_name = "waynebruce2110/GraveSocialAI"
 
-# Enable 8-bit quantization for CPU
-quantization_config = BitsAndBytesConfig(
-    load_in_8bit=True,  # Enables 8-bit loading
-    llm_int8_enable_fp32_cpu_offload=True  # Ensures it works on CPU
-)
-
-# Load the tokenizer and model with 8-bit quantization
+# Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained(model_name, local_files_only=False)
+
+# Load model with 8-bit quantization on CPU
 model = AutoModelForCausalLM.from_pretrained(
-    model_name, 
-    local_files_only=False, 
-    device_map="cpu",  # Ensures it loads on CPU
-    quantization_config=quantization_config
+    model_name,
+    local_files_only=False,
+    torch_dtype=torch.float16,  # Reduces memory usage
+    device_map="cpu"  # Forces model to load on CPU
 )
 
 # Define input schema
@@ -34,8 +30,9 @@ def read_root():
 
 @app.post("/generate/")
 def generate_text(data: PromptInput):
-    inputs = tokenizer(data.prompt, return_tensors="pt")
-    outputs = model.generate(**inputs, max_length=100)
+    inputs = tokenizer(data.prompt, return_tensors="pt").to("cpu")  # Ensure input is on CPU
+    with torch.no_grad():
+        outputs = model.generate(**inputs, max_length=100)
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return {"generated_text": response}