tcy6 commited on
Commit
76a6782
1 Parent(s): dc55a3a

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +6 -6
app.py CHANGED
@@ -19,7 +19,7 @@ import json
19
  cache_dir = '/data/KB'
20
  os.makedirs(cache_dir, exist_ok=True)
21
 
22
- @spaces.GPU(duration=100)
23
  def weighted_mean_pooling(hidden, attention_mask):
24
  attention_mask_ = attention_mask * attention_mask.cumsum(dim=1)
25
  s = torch.sum(hidden * attention_mask_.unsqueeze(-1).float(), dim=1)
@@ -27,7 +27,7 @@ def weighted_mean_pooling(hidden, attention_mask):
27
  reps = s / d
28
  return reps
29
 
30
- @spaces.GPU(duration=100)
31
  @torch.no_grad()
32
  def encode(text_or_image_list):
33
  global model, tokenizer
@@ -63,7 +63,7 @@ def calculate_md5_from_binary(binary_data):
63
  hash_md5.update(binary_data)
64
  return hash_md5.hexdigest()
65
 
66
- @spaces.GPU(duration=100)
67
  def add_pdf_gradio(pdf_file_binary, progress=gr.Progress()):
68
  global model, tokenizer
69
  model.eval()
@@ -108,7 +108,7 @@ def add_pdf_gradio(pdf_file_binary, progress=gr.Progress()):
108
 
109
  return knowledge_base_name
110
 
111
- @spaces.GPU(duration=100)
112
  def retrieve_gradio(knowledge_base: str, query: str, topk: int):
113
  global model, tokenizer
114
 
@@ -224,7 +224,7 @@ gen_model.to(device)
224
  print("gen model load success!")
225
 
226
 
227
- @spaces.GPU(duration=50)
228
  def answer_question(images, question):
229
  global gen_model, gen_tokenizer
230
  # here each element of images is a tuple of (image_path, None).
@@ -257,7 +257,7 @@ with gr.Blocks() as app:
257
  gr.Markdown("- Currently online demo support PDF document with less than 50 pages due to GPU time limit. Deploy on your own machine for longer PDFs and books.")
258
 
259
  with gr.Row():
260
- file_input = gr.File(type="binary", label="Step 1: Upload PDF")
261
  file_result = gr.Text(label="Knowledge Base ID (remember it, it is re-usable!)")
262
  process_button = gr.Button("Process PDF (Don't click until PDF upload success)")
263
 
 
19
  cache_dir = '/data/KB'
20
  os.makedirs(cache_dir, exist_ok=True)
21
 
22
+ @spaces.GPU
23
  def weighted_mean_pooling(hidden, attention_mask):
24
  attention_mask_ = attention_mask * attention_mask.cumsum(dim=1)
25
  s = torch.sum(hidden * attention_mask_.unsqueeze(-1).float(), dim=1)
 
27
  reps = s / d
28
  return reps
29
 
30
+ @spaces.GPU
31
  @torch.no_grad()
32
  def encode(text_or_image_list):
33
  global model, tokenizer
 
63
  hash_md5.update(binary_data)
64
  return hash_md5.hexdigest()
65
 
66
+ @spaces.GPU
67
  def add_pdf_gradio(pdf_file_binary, progress=gr.Progress()):
68
  global model, tokenizer
69
  model.eval()
 
108
 
109
  return knowledge_base_name
110
 
111
+ @spaces.GPU
112
  def retrieve_gradio(knowledge_base: str, query: str, topk: int):
113
  global model, tokenizer
114
 
 
224
  print("gen model load success!")
225
 
226
 
227
+ @spaces.GPU
228
  def answer_question(images, question):
229
  global gen_model, gen_tokenizer
230
  # here each element of images is a tuple of (image_path, None).
 
257
  gr.Markdown("- Currently online demo support PDF document with less than 50 pages due to GPU time limit. Deploy on your own machine for longer PDFs and books.")
258
 
259
  with gr.Row():
260
+ file_input = gr.File(file_types=["pdf"], file_count="multiple", type="binary", label="Step 1: Upload PDF")
261
  file_result = gr.Text(label="Knowledge Base ID (remember it, it is re-usable!)")
262
  process_button = gr.Button("Process PDF (Don't click until PDF upload success)")
263