Upload app.py

app.py

@@ -19,7 +19,7 @@ import json
 cache_dir = '/data/KB'
 os.makedirs(cache_dir, exist_ok=True)
 
-@spaces.GPU
+@spaces.GPU(duration=100)
 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)
@@ -27,7 +27,7 @@ def weighted_mean_pooling(hidden, attention_mask):
     reps = s / d
     return reps
 
-@spaces.GPU
+@spaces.GPU(duration=100)
 @torch.no_grad()
 def encode(text_or_image_list):
     global model, tokenizer
@@ -108,7 +108,7 @@ def add_pdf_gradio(pdf_file_binary, progress=gr.Progress()):
     
     return knowledge_base_name
 
-@spaces.GPU
+@spaces.GPU(duration=100)
 def retrieve_gradio(knowledge_base: str, query: str, topk: int):
     global model, tokenizer