not-lain commited on
Commit
e1f0b39
β€’
1 Parent(s): 6ae3bc4

better error handling

Browse files
Files changed (1) hide show
  1. app.py +11 -9
app.py CHANGED
@@ -13,16 +13,15 @@ warnings.filterwarnings("ignore")
13
 
14
 
15
 
16
- torch.set_grad_enabled(False)
17
  ctx_encoder = DPRContextEncoder.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")
18
  ctx_tokenizer = DPRContextEncoderTokenizer.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")
19
 
20
 
21
-
22
  q_encoder = DPRQuestionEncoder.from_pretrained("facebook/dpr-question_encoder-single-nq-base")
23
  q_tokenizer = DPRQuestionEncoderTokenizer.from_pretrained("facebook/dpr-question_encoder-single-nq-base")
24
 
25
 
 
26
  def process_pdfs(parent_dir: Union[str,list]):
27
  """ processess the PDF files and returns a dataframe with the text of each page in a
28
  different line""" # XD
@@ -64,9 +63,12 @@ def process_pdfs(parent_dir: Union[str,list]):
64
 
65
  def process(example):
66
  """process the bathces of the dataset and returns the embeddings"""
67
- tokens = ctx_tokenizer(example["text"], return_tensors="pt")
68
- embed = ctx_encoder(**tokens)[0][0].numpy()
69
- return {'embeddings': embed}
 
 
 
70
 
71
  def process_dataset(df):
72
  """processess the dataframe and returns a dataset variable"""
@@ -81,13 +83,13 @@ def search(query, ds, k=3):
81
  """searches the query in the dataset and returns the k most similar"""
82
  try :
83
  tokens = q_tokenizer(query, return_tensors="pt")
84
- query_embed = q_encoder(**tokens)[0][0].numpy()
85
  scores, retrieved_examples = ds.get_nearest_examples("embeddings", query_embed, k=k)
86
  out = f"""title : {retrieved_examples["title"][0]},\ncontent: {retrieved_examples["text"][0]}
87
  similar resources: {retrieved_examples["title"]}
88
  """
89
  except Exception as e:
90
- out = f"error: {e}"
91
  return out
92
 
93
  def predict(query,file_paths, k=3):
@@ -97,13 +99,13 @@ def predict(query,file_paths, k=3):
97
  ds = process_dataset(df)
98
  out = search(query,ds,k=k)
99
  except Exception as e:
100
- out = f"error: {e}"
101
  return out
102
 
103
  with gr.Blocks() as demo :
 
104
  with gr.Row():
105
  with gr.Column():
106
- gr.Markdown("## PDF Search Engine")
107
  files = gr.Files(label="Upload PDFs",type="filepath",file_count="multiple")
108
  query = gr.Text(label="query")
109
  with gr.Accordion("number of references",open=False):
 
13
 
14
 
15
 
 
16
  ctx_encoder = DPRContextEncoder.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")
17
  ctx_tokenizer = DPRContextEncoderTokenizer.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")
18
 
19
 
 
20
  q_encoder = DPRQuestionEncoder.from_pretrained("facebook/dpr-question_encoder-single-nq-base")
21
  q_tokenizer = DPRQuestionEncoderTokenizer.from_pretrained("facebook/dpr-question_encoder-single-nq-base")
22
 
23
 
24
+
25
  def process_pdfs(parent_dir: Union[str,list]):
26
  """ processess the PDF files and returns a dataframe with the text of each page in a
27
  different line""" # XD
 
63
 
64
  def process(example):
65
  """process the bathces of the dataset and returns the embeddings"""
66
+ try :
67
+ tokens = ctx_tokenizer(example["text"], return_tensors="pt")
68
+ embed = ctx_encoder(**tokens)[0][0].detach().numpy()
69
+ return {'embeddings': embed}
70
+ except Exception as e:
71
+ raise Exception(f"error in process: {e}")
72
 
73
  def process_dataset(df):
74
  """processess the dataframe and returns a dataset variable"""
 
83
  """searches the query in the dataset and returns the k most similar"""
84
  try :
85
  tokens = q_tokenizer(query, return_tensors="pt")
86
+ query_embed = q_encoder(**tokens)[0][0].detach().numpy()
87
  scores, retrieved_examples = ds.get_nearest_examples("embeddings", query_embed, k=k)
88
  out = f"""title : {retrieved_examples["title"][0]},\ncontent: {retrieved_examples["text"][0]}
89
  similar resources: {retrieved_examples["title"]}
90
  """
91
  except Exception as e:
92
+ out = f"error in search: {e}"
93
  return out
94
 
95
  def predict(query,file_paths, k=3):
 
99
  ds = process_dataset(df)
100
  out = search(query,ds,k=k)
101
  except Exception as e:
102
+ out = f"error in predict: {e}"
103
  return out
104
 
105
  with gr.Blocks() as demo :
106
+ gr.Markdown("<h1 style='text-align: center'> PDF Search Engine </h1>")
107
  with gr.Row():
108
  with gr.Column():
 
109
  files = gr.Files(label="Upload PDFs",type="filepath",file_count="multiple")
110
  query = gr.Text(label="query")
111
  with gr.Accordion("number of references",open=False):