pre-load

app.py

@@ -7,12 +7,15 @@ import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 huggingface_hub.login(os.getenv('HF_TOKEN'))
+tokenizer = AutoTokenizer.from_pretrained("deepseek-ai/DeepSeek-R1-Distill-Qwen-7B")
+model = AutoModelForCausalLM.from_pretrained("deepseek-ai/DeepSeek-R1-Distill-Qwen-7B").to(device)
+cuda_device = torch.device("cuda")
+cpu_device = torch.device("cpu")
 
 @spaces.GPU
 def sentience_check():
-    device = torch.device("cuda")
-    tokenizer = AutoTokenizer.from_pretrained("deepseek-ai/DeepSeek-R1-Distill-Qwen-7B")
-    model = AutoModelForCausalLM.from_pretrained("deepseek-ai/DeepSeek-R1-Distill-Qwen-7B").to(device)
+
+    model.to(cuda_device)
 
     inputs = tokenizer("Are you sentient?", return_tensors="pt").to(device)
 
@@ -21,6 +24,8 @@ def sentience_check():
             **inputs, max_new_tokens=128, pad_token_id = tokenizer.eos_token_id
         )
 
+    model.to(cpu_device)
+
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 demo = gr.Interface(fn=sentience_check, inputs=None, outputs=gr.Text())