Update app.py

app.py

@@ -1,3 +1,4 @@
+import spaces
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
@@ -7,22 +8,20 @@ import requests
 from openai import OpenAI
 
 HF_API_KEY = os.getenv("HF_API_KEY")
-model_id = "meta-llama/Llama-Guard-3-1B"
-
+MODEL_ID = "meta-llama/Llama-Guard-3-1B"
 
 model = AutoModelForCausalLM.from_pretrained(
-    model_id,
+    MODEL_ID,
     torch_dtype=torch.bfloat16,
     device_map="auto",
     token=HF_API_KEY
-    
-)
-tokenizer = AutoTokenizer.from_pretrained(model_id, token=HF_API_KEY)
+).to('cuda')
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_API_KEY).to('cuda')
 
 
 # Model 1: Llama Guard Model
+@spaces.GPU
 def llama_guard_moderation(input_text):
-
     conversation = [
         {
             "role": "user",
@@ -34,7 +33,6 @@ def llama_guard_moderation(input_text):
             ],
         }
     ]
-
     input_ids = tokenizer.apply_chat_template(
         conversation, return_tensors="pt"
     ).to(model.device)