Update app.py

app.py

@@ -3,6 +3,7 @@ import torch
 from flask import Flask, jsonify, request
 from flask_cors import CORS
 from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, GenerationConfig
+from accelerate import Accelerator
 import re
 import traceback
 
@@ -17,6 +18,7 @@ CORS(app, resources={r"/send_message": {"origins": ["http://localhost:3000", "ht
 # Global variables for model and tokenizer
 model = None
 tokenizer = None
+accelerator = Accelerator()
 
 def get_model_and_tokenizer(model_id: str):
     global model, tokenizer
@@ -36,11 +38,14 @@ def get_model_and_tokenizer(model_id: str):
             model = AutoModelForCausalLM.from_pretrained(
                 model_id, quantization_config=bnb_config, device_map="auto"
             )
-            
+
             model.config.use_cache = False
             model.config.pretraining_tp = 1
             model.config.pad_token_id = tokenizer.eos_token_id  # Fix padding issue
 
+            # Ensure model is placed on the correct device using accelerate
+            model = accelerator.prepare(model)
+
         except Exception as e:
             print("Error loading model:")
             print(traceback.format_exc())  # Logs the full error traceback
@@ -51,7 +56,7 @@ def generate_response(user_input, model_id):
         get_model_and_tokenizer(model_id)
 
         prompt = user_input
-        device = "cuda" if torch.cuda.is_available() else "cpu"
+        device = accelerator.device  # Automatically uses GPU or CPU based on accelerator setup
 
         generation_config = GenerationConfig(
             penalty_alpha=0.6,
@@ -103,4 +108,4 @@ def handle_post_request():
         return jsonify({"error": str(e)}), 500
 
 if __name__ == '__main__':
-    app.run(host='0.0.0.0', port=7860)
+    app.run(host='0.0.0.0', port=7860)