better error handling

app.py

@@ -13,16 +13,15 @@ warnings.filterwarnings("ignore")
 
 
 
-torch.set_grad_enabled(False)
 ctx_encoder = DPRContextEncoder.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")
 ctx_tokenizer = DPRContextEncoderTokenizer.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")
 
 
-
 q_encoder = DPRQuestionEncoder.from_pretrained("facebook/dpr-question_encoder-single-nq-base")
 q_tokenizer = DPRQuestionEncoderTokenizer.from_pretrained("facebook/dpr-question_encoder-single-nq-base")
 
 
+
 def process_pdfs(parent_dir: Union[str,list]):
     """ processess the PDF files and returns a dataframe with the text of each page in a
     different line"""  # XD
@@ -64,9 +63,12 @@ def process_pdfs(parent_dir: Union[str,list]):
 
 def process(example):
   """process the bathces of the dataset and returns the embeddings"""
-  tokens = ctx_tokenizer(example["text"], return_tensors="pt")
-  embed = ctx_encoder(**tokens)[0][0].numpy()
-  return {'embeddings': embed}
+  try : 
+    tokens = ctx_tokenizer(example["text"], return_tensors="pt")
+    embed = ctx_encoder(**tokens)[0][0].detach().numpy()
+    return {'embeddings': embed}
+  except Exception as e:
+    raise Exception(f"error in process: {e}")
 
 def process_dataset(df):
     """processess the dataframe and returns a dataset variable"""
@@ -81,13 +83,13 @@ def search(query, ds, k=3):
     """searches the query in the dataset and returns the k most similar"""
     try :
         tokens = q_tokenizer(query, return_tensors="pt")
-        query_embed = q_encoder(**tokens)[0][0].numpy()
+        query_embed = q_encoder(**tokens)[0][0].detach().numpy()
         scores, retrieved_examples = ds.get_nearest_examples("embeddings", query_embed, k=k)
         out = f"""title : {retrieved_examples["title"][0]},\ncontent: {retrieved_examples["text"][0]}
         similar resources: {retrieved_examples["title"]}
         """
     except Exception as e:
-        out = f"error: {e}"
+        out = f"error in search: {e}"
     return out
     
 def predict(query,file_paths, k=3):
@@ -97,13 +99,13 @@ def predict(query,file_paths, k=3):
         ds = process_dataset(df)
         out =  search(query,ds,k=k)
     except Exception as e:
-        out = f"error: {e}"
+        out = f"error in predict: {e}"
     return out
 
 with gr.Blocks() as demo : 
+    gr.Markdown("<h1 style='text-align: center'> PDF Search Engine </h1>")
     with gr.Row():
         with gr.Column(): 
-            gr.Markdown("## PDF Search Engine")
             files = gr.Files(label="Upload PDFs",type="filepath",file_count="multiple")
             query = gr.Text(label="query")
             with gr.Accordion("number of references",open=False):