Update app.py

app.py

@@ -19,6 +19,7 @@ import json
 cache_dir = '/data/KB'
 os.makedirs(cache_dir, exist_ok=True)
 
+@spaces.GPU
 def weighted_mean_pooling(hidden, attention_mask):
     attention_mask_ = attention_mask * attention_mask.cumsum(dim=1)
     s = torch.sum(hidden * attention_mask_.unsqueeze(-1).float(), dim=1)
@@ -26,6 +27,7 @@ def weighted_mean_pooling(hidden, attention_mask):
     reps = s / d
     return reps
 
+@spaces.GPU
 @torch.no_grad()
 def encode(text_or_image_list):
     global model, tokenizer
@@ -106,7 +108,7 @@ def add_pdf_gradio(pdf_file_binary, progress=gr.Progress()):
     
     return knowledge_base_name
 
-# @spaces.GPU
+@spaces.GPU
 def retrieve_gradio(knowledge_base: str, query: str, topk: int):
     global model, tokenizer