ee

app.py

@@ -1,4 +1,6 @@
 import gradio as gr
+import spaces
+
 # from huggingface_hub import InferenceClient
 from transformers import pipeline
 from transformers import AutoModelForCausalLM, AutoTokenizer
@@ -8,8 +10,9 @@ HF_TOKEN = os.getenv('HF_TOKEN')
 checkpoint = "zidsi/SLlamica_PT4SFT_v2"
 device = "cuda"  # "cuda" or "cpu"
 tokenizer = AutoTokenizer.from_pretrained(checkpoint,token=HF_TOKEN)
-model = AutoModelForCausalLM.from_pretrained(checkpoint,token=HF_TOKEN).to(device)
-
+model = AutoModelForCausalLM.from_pretrained(checkpoint,token=HF_TOKEN)
+model.to(device)
+@spaces.GPU
 def predict(message, history):
     history.append({"role": "user", "content": message})
     input_text = tokenizer.apply_chat_template(history, tokenize=False)