device option cuda or cpu

app.py

@@ -1,4 +1,5 @@
 import os
+import torch
 from flask import Flask, jsonify, request
 from flask_cors import CORS 
 from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig  
@@ -14,6 +15,7 @@ CORS(app, resources={r"api/predict/*": {"origins": ["http://localhost:3000", "ht
 # Global variables for model and tokenizer
 model = None
 tokenizer = None
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  # Use GPU if available
 
 def get_model_and_tokenizer(model_id):
     global model, tokenizer
@@ -23,9 +25,9 @@ def get_model_and_tokenizer(model_id):
         tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=False)
         tokenizer.pad_token = tokenizer.eos_token
         
-        print(f"Loading model and for model_id: {model_id}")
-        # Load the model
-        model = AutoModelForCausalLM.from_pretrained(model_id) #, device_map="auto")
+        print(f"Loading model for model_id: {model_id}")
+        # Load the model and move it to the specified device
+        model = AutoModelForCausalLM.from_pretrained(model_id).to(device) #, device_map="auto")
         model.config.use_cache = False
    
     except Exception as e:
@@ -38,8 +40,8 @@ def generate_response(user_input, model_id):
     if model is None or tokenizer is None:
         get_model_and_tokenizer(model_id)  # Load model and tokenizer
 
-    # Prepare the input tensors
-    inputs = tokenizer(prompt, return_tensors="pt")  # Move inputs to GPU if available
+    # Prepare the input tensors and move them to the appropriate device
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
     
     generation_config = GenerationConfig(
         max_new_tokens=100,

requirements.txt

@@ -1,9 +1,9 @@
 flask
 flask_cors
 huggingface-hub
-transformers
-torch
-accelerate
-bitsandbytes
+transformers>=4.30.0
+torch>=2.0.0
+accelerate>=0.18.0
+bitsandbytes-cuda117  # Replace 'cuda117' with your specific CUDA version if different
 peft
-trl
+trl