use gpu only under decorated spaces.GPU

app.py

@@ -68,7 +68,15 @@ MAX_PROMPT_TOKENS = 256
 
 
 @spaces.GPU
-def completion(prompt: str, model, tokenizer):
+def completion(prompt: str, model_id: str):
+    # load model and tokenizer
+    logger.info(f"loading local model and tokenizer for {model_id}")
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+
+    dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float16
+    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=dtype, device_map="auto")
+    logger.info(f"done loading {model_id}")
+
     # tokenize
     input_ids = tokenizer.apply_chat_template(
         [
@@ -93,6 +101,12 @@ def completion(prompt: str, model, tokenizer):
         top_p=None,
         temperature=None,
     )
+
+    # cleanup memory
+    del model, tokenizer
+    torch.cuda.empty_cache()
+    gc.collect()
+
     return tokenizer.decode(outputs[0][input_ids.shape[-1] :], skip_special_tokens=True)
 
 
@@ -107,16 +121,6 @@ def completion_openrouter(prompt: str, model_id: str):
     return resp.choices[0].message.content
 
 
-# @functools.cache
-def load_model_and_tokenizer(model_id: str):
-    logger.info(f"loading local model and tokenizer for {model_id}")
-    tokenizer = AutoTokenizer.from_pretrained(model_id)
-    dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float16
-    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=dtype, device_map="auto")
-    logger.info(f"done loading {model_id}")
-    return model, tokenizer
-
-
 def load_openrouter_client():
     logger.info(f"connecting to OpenRouter")
     return OpenAI(
@@ -135,13 +139,7 @@ def get_completion(*, prompt: str, model_id: str):
     if model_id.startswith("api:"):
         return completion_openrouter(prompt, model_id.removeprefix("api:"))
     else:
-        model, tokenizer = load_model_and_tokenizer(model_id)
-        resp = completion(prompt, model, tokenizer)
-
-        # cleanup memory
-        del model, tokenizer
-        torch.cuda.empty_cache()
-        gc.collect()
+        resp = completion(prompt, model_id)
         return resp