Update app.py

app.py

@@ -16,7 +16,6 @@ app = FastAPI(root_path="/api/v1")
 
 # Load the model and tokenizer
 model_name_or_path = "TheBloke/Wizard-Vicuna-7B-Uncensored-GPTQ"
-mistral_model="mistralai/Mistral-7B-Instruct-v0.2"
 
 conversations = {}
 device = "cuda" # the device to load the model onto
@@ -29,12 +28,9 @@ def mistral_model():
         tuple: A tuple containing the loaded model and tokenizer.
     """
     if torch.cuda.is_available():
-        print("CUDA is available. GPU will be used.")
-    else:
-        print("CUDA is not available. CPU will be used.")
-   
-    tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2")
-    model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2")
+        model_id = "mistralai/Mistral-7B-Instruct-v0.2"
+        model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map="auto")
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
 
     return model,tokenizer