Spaces:
Running
Running
| import os | |
| import PyPDF2 | |
| import pandas as pd | |
| import warnings | |
| import re | |
| from transformers import DPRContextEncoder, DPRContextEncoderTokenizer | |
| from transformers import DPRQuestionEncoder, DPRQuestionEncoderTokenizer | |
| import torch | |
| import gradio as gr | |
| from typing import Union | |
| from datasets import Dataset | |
| warnings.filterwarnings("ignore") | |
| 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 | |
| # creating a pdf file object | |
| df = pd.DataFrame(columns = ["title","text"]) | |
| if type(parent_dir) == str : | |
| parent_dir = [parent_dir] | |
| for file_path in parent_dir: | |
| if ".pdf" not in file_path : # skip non pdf files | |
| raise Exception("only pdf files are supported") | |
| # creating a pdf file object | |
| pdfFileObj = open(file_path, 'rb') | |
| # creating a pdf reader object | |
| pdfReader = PyPDF2.PdfReader(pdfFileObj) | |
| # printing number of pages in pdf file | |
| num_pages = len(pdfReader.pages) | |
| for i in range(num_pages) : | |
| pageObj = pdfReader.pages[i] | |
| # extracting text from page | |
| txt = pageObj.extract_text() | |
| txt = txt.replace("\n","") # strip return to line | |
| txt = txt.replace("\t","") # strip tabs | |
| txt = re.sub(r" +"," ",txt) # strip extra space | |
| # 512 is related to the positional encoding "facebook/dpr-ctx_encoder-single-nq-base" model | |
| file_name = file_path.split("/")[-1] | |
| if len(txt) < 512 : | |
| new_data = pd.DataFrame([[f"{file_name}-page-{i}",txt]],columns=["title","text"]) | |
| df = pd.concat([df,new_data],ignore_index=True) | |
| else : | |
| while len(txt) > 512 : | |
| new_data = pd.DataFrame([[f"{file_name}-page-{i}",txt[:512]]],columns=["title","text"]) | |
| df = pd.concat([df,new_data],ignore_index=True) | |
| txt = txt[512:] | |
| # closing the pdf file object | |
| pdfFileObj.close() | |
| return df | |
| def process(example): | |
| """process the bathces of the dataset and returns the embeddings""" | |
| 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""" | |
| if len(df) == 0 : | |
| raise Exception("empty pdf files, or can't read text from them") | |
| ds = Dataset.from_pandas(df) | |
| ds = ds.map(process) | |
| ds.add_faiss_index(column='embeddings') # add faiss index | |
| return ds | |
| 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].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]}\n\n\n**similar resources:** {retrieved_examples["title"]} | |
| """ | |
| except Exception as e: | |
| out = f"error in search: {e}" | |
| return out | |
| def predict(query,file_paths, k=3): | |
| """predicts the most similar files to the query""" | |
| try : | |
| df = process_pdfs(file_paths) | |
| ds = process_dataset(df) | |
| out = search(query,ds,k=k) | |
| except Exception as 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(): | |
| files = gr.Files(label="Upload PDFs",type="filepath",file_count="multiple") | |
| query = gr.Text(label="query") | |
| with gr.Accordion("number of references",open=False): | |
| k = gr.Number(value=3,show_label=False,precision=0,minimum=1,container=False) | |
| button = gr.Button("search") | |
| with gr.Column(): | |
| output = gr.Markdown(label="output") | |
| button.click(predict, [query,files,k],outputs=output) | |
| demo.launch() | |