Upload app.py

app.py

@@ -1,5 +1,5 @@
 # Install the necessary packages
-# pip install accelerate transformers fastapi pydantic torch
+# pip install accelerate transformers fastapi pydantic torch jinja2
 
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
@@ -9,14 +9,15 @@ from fastapi import FastAPI
 # Initialize the FastAPI app
 app = FastAPI(docs_url="/")
 
-# Load the model and tokenizer once at startup
-device = "cuda" # the device to load the model onto
+# Determine the device to use
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
+# Load the model and tokenizer once at startup
 model = AutoModelForCausalLM.from_pretrained(
     "Qwen/Qwen1.5-0.5B-Chat",
     torch_dtype="auto",
     device_map="auto"
-)
+).to(device)
 tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen1.5-0.5B-Chat")
 
 # Define the request model