Spaces:
Build error
Build error
| from pathlib import Path | |
| from arxiv_public_data.fulltext import convert_directory_parallel | |
| from arxiv_public_data import internal_citations | |
| import torch | |
| import os | |
| from summarizer import Summarizer | |
| from sentence_transformers import SentenceTransformer | |
| import spacy | |
| import numpy as np | |
| from keybert import KeyBERT | |
| import shutil, joblib | |
| from distutils.dir_util import copy_tree | |
| try: | |
| from transformers import * | |
| except: | |
| from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoConfig, AutoModel, LEDTokenizer, \ | |
| LEDForConditionalGeneration | |
| from src.defaults import DEFAULTS_CPU_COMPAT, DEFAULTS_HIGH_GPU | |
| class Surveyor: | |
| ''' | |
| A class to abstract all nlp and data mining helper functions as well as workflows | |
| required to generate the survey from a single query, with absolute configurability | |
| ''' | |
| def __init__( | |
| self, | |
| models_dir=None, | |
| title_model_name=None, | |
| ex_summ_model_name=None, | |
| ledmodel_name=None, | |
| embedder_name=None, | |
| nlp_name=None, | |
| similarity_nlp_name=None, | |
| kw_model_name=None, | |
| high_gpu=False, | |
| refresh_models=False, | |
| no_save_models=False, | |
| print_fn=None, | |
| survey_print_fn=None | |
| ): | |
| ''' | |
| Initializes models and directory structure for the surveyor | |
| Optional Params: | |
| - models_dir: String, directory to save to huge models | |
| - title_model_name: String, title model name/tag in hugging-face, defaults to `Callidior/bert2bert-base-arxiv-titlegen` | |
| - ex_summ_model_name: String, extractive summary model name/tag in hugging-face, defaults to `allenai/scibert_scivocab_uncased` | |
| - ledmodel_name: String, led model(for abstractive summary) name/tag in hugging-face, defaults to `allenai/led-large-16384-arxiv` | |
| - embedder_name: String, sentence embedder name/tag in hugging-face, defaults to `paraphrase-MiniLM-L6-v2` | |
| - nlp_name: String, spacy model name/tag in hugging-face (if changed - needs to be spacy-installed prior), defaults to `en_core_sci_scibert` | |
| - similarity_nlp_name: String, spacy downstream trained model(for similarity) name/tag in hugging-face (if changed - needs to be spacy-installed prior), defaults to `en_core_sci_lg` | |
| - kw_model_name: String, keyword extraction model name/tag in hugging-face, defaults to `distilbert-base-nli-mean-tokens` | |
| - high_gpu: Bool, High GPU usage permitted, defaults to False | |
| - refresh_models: Bool, Refresh model downloads with given names (needs atleast one model name param above), defaults to False | |
| - no_save_models: forces refresh models | |
| - max_search: int maximium number of papers to gaze at - defaults to 100 | |
| - num_papers: int maximium number of papers to download and analyse - defaults to 25 | |
| ''' | |
| self.print_fn = print | |
| if print_fn is not None: | |
| self.print_fn = print_fn | |
| self.survey_print_fn = self.print_fn | |
| if survey_print_fn is not None: | |
| self.survey_print_fn = survey_print_fn | |
| self.torch_device = 'cpu' | |
| self.print_fn("\nTorch_device: " + self.torch_device) | |
| if torch.cuda.is_available(): | |
| self.torch_device = 'cuda' | |
| spacy.require_gpu() | |
| self.high_gpu = high_gpu | |
| self.DEFAULTS = DEFAULTS_CPU_COMPAT | |
| if self.high_gpu: | |
| self.DEFAULTS = DEFAULTS_HIGH_GPU | |
| if not kw_model_name: | |
| kw_model_name = self.DEFAULTS["kw_model_name"] | |
| self.num_papers = self.DEFAULTS['num_papers'] | |
| self.max_search = self.DEFAULTS['max_search'] | |
| if not models_dir: | |
| models_dir = self.DEFAULTS['models_dir'] | |
| models_found = False | |
| if os.path.exists(models_dir) and not no_save_models: | |
| if len(os.listdir(models_dir)) > 3: | |
| models_found = True | |
| if not title_model_name: | |
| title_model_name = self.DEFAULTS["title_model_name"] | |
| if not ex_summ_model_name: | |
| ex_summ_model_name = self.DEFAULTS["ex_summ_model_name"] | |
| if not ledmodel_name: | |
| ledmodel_name = self.DEFAULTS["ledmodel_name"] | |
| if not embedder_name: | |
| embedder_name = self.DEFAULTS["embedder_name"] | |
| if not nlp_name: | |
| nlp_name = self.DEFAULTS["nlp_name"] | |
| if not similarity_nlp_name: | |
| similarity_nlp_name = self.DEFAULTS["similarity_nlp_name"] | |
| if refresh_models or not models_found: | |
| self.print_fn(f'\nInitializing models {"and saving (about 5GB)" if not no_save_models else ""}') | |
| if not no_save_models: | |
| self.clean_dirs([models_dir]) | |
| self.title_tokenizer = AutoTokenizer.from_pretrained(title_model_name, trust_remote_code=True) | |
| self.title_model = AutoModelForSeq2SeqLM.from_pretrained(title_model_name, trust_remote_code=True).to(self.torch_device) | |
| self.title_model.eval() | |
| if not no_save_models: | |
| self.title_model.save_pretrained(models_dir + "/title_model") | |
| #self.title_tokenizer.save_pretrained(models_dir + "/title_tokenizer") | |
| # summary model | |
| self.custom_config = AutoConfig.from_pretrained(ex_summ_model_name) | |
| self.custom_config.output_hidden_states = True | |
| self.summ_tokenizer = AutoTokenizer.from_pretrained(ex_summ_model_name) | |
| self.summ_model = AutoModel.from_pretrained(ex_summ_model_name, config=self.custom_config).to( | |
| self.torch_device) | |
| self.summ_model.eval() | |
| if not no_save_models: | |
| self.summ_model.save_pretrained(models_dir + "/summ_model") | |
| #self.summ_tokenizer.save_pretrained(models_dir + "/summ_tokenizer") | |
| self.model = Summarizer(custom_model=self.summ_model, custom_tokenizer=self.summ_tokenizer) | |
| if 'led' in ledmodel_name: | |
| self.ledtokenizer = LEDTokenizer.from_pretrained(ledmodel_name) | |
| self.ledmodel = LEDForConditionalGeneration.from_pretrained(ledmodel_name).to(self.torch_device) | |
| elif 't5' in ledmodel_name: | |
| self.ledtokenizer = AutoTokenizer.from_pretrained(ledmodel_name) | |
| self.ledmodel = T5ForConditionalGeneration.from_pretrained(ledmodel_name).to(self.torch_device) | |
| elif 'bart' in ledmodel_name: | |
| self.ledtokenizer = AutoTokenizer.from_pretrained(ledmodel_name) | |
| self.ledmodel = BartForConditionalGeneration.from_pretrained(ledmodel_name).to(self.torch_device) | |
| self.ledmodel.eval() | |
| if not no_save_models: | |
| self.ledmodel.save_pretrained(models_dir + "/ledmodel") | |
| #self.ledtokenizer.save_pretrained(models_dir + "/ledtokenizer") | |
| self.embedder = SentenceTransformer(embedder_name) | |
| self.embedder.eval() | |
| if not no_save_models: | |
| self.embedder.save(models_dir + "/embedder") | |
| else: | |
| self.print_fn("\n- Initializing from previously saved models at" + models_dir) | |
| self.title_tokenizer = AutoTokenizer.from_pretrained(title_model_name) | |
| self.title_model = AutoModelForSeq2SeqLM.from_pretrained(models_dir + "/title_model").to(self.torch_device) | |
| self.title_model.eval() | |
| # summary model | |
| #self.summ_config = AutoConfig.from_pretrained(ex_summ_model_name) | |
| #self.summ_config.output_hidden_states = True | |
| self.summ_tokenizer = AutoTokenizer.from_pretrained(ex_summ_model_name) | |
| self.summ_model = AutoModel.from_pretrained(models_dir + "/summ_model").to( | |
| self.torch_device) | |
| self.summ_model.eval() | |
| self.model = Summarizer(custom_model=self.summ_model, custom_tokenizer=self.summ_tokenizer) | |
| if 'led' in ledmodel_name: | |
| self.ledtokenizer = LEDTokenizer.from_pretrained(ledmodel_name) | |
| self.ledmodel = LEDForConditionalGeneration.from_pretrained(models_dir + "/ledmodel").to(self.torch_device) | |
| elif 't5' in ledmodel_name: | |
| self.ledtokenizer = AutoTokenizer.from_pretrained(ledmodel_name) | |
| self.ledmodel = T5ForConditionalGeneration.from_pretrained(models_dir + "/ledmodel").to(self.torch_device) | |
| elif 'bart' in ledmodel_name: | |
| self.ledtokenizer = AutoTokenizer.from_pretrained(ledmodel_name) | |
| self.ledmodel = BartForConditionalGeneration.from_pretrained(models_dir + "/ledmodel").to(self.torch_device) | |
| self.ledmodel.eval() | |
| self.embedder = SentenceTransformer(models_dir + "/embedder") | |
| self.embedder.eval() | |
| self.nlp = spacy.load(nlp_name) | |
| self.similarity_nlp = spacy.load(similarity_nlp_name) | |
| self.kw_model = KeyBERT(kw_model_name) | |
| def define_structure(self, pdf_dir=None, txt_dir=None, img_dir=None, tab_dir=None, dump_dir=None): | |
| if pdf_dir: | |
| survey_pdf_dir = pdf_dir | |
| else: | |
| survey_pdf_dir = self.DEFAULTS["pdf_dir"] | |
| if txt_dir: | |
| survey_txt_dir = txt_dir | |
| else: | |
| survey_txt_dir = self.DEFAULTS["txt_dir"] | |
| if img_dir: | |
| survey_img_dir = img_dir | |
| else: | |
| survey_img_dir = self.DEFAULTS["img_dir"] | |
| if tab_dir: | |
| survey_tab_dir = tab_dir | |
| else: | |
| survey_tab_dir = self.DEFAULTS["tab_dir"] | |
| if dump_dir: | |
| survey_dump_dir = dump_dir | |
| else: | |
| survey_dump_dir = self.DEFAULTS["dump_dir"] | |
| dirs = [survey_pdf_dir, survey_txt_dir, survey_img_dir, survey_tab_dir, survey_dump_dir] | |
| if sum([True for dir in dirs if 'arxiv_data/' in dir]): | |
| base = os.path.dirname("arxiv_data/") | |
| if not os.path.exists(base): | |
| os.mkdir(base) | |
| self.clean_dirs(dirs) | |
| return dirs | |
| def clean_dirs(self, dirs): | |
| import shutil | |
| for d in dirs: | |
| if os.path.exists(d): | |
| shutil.rmtree(d) | |
| os.mkdir(d) | |
| def pdf_route(self, pdf_dir, txt_dir, img_dir, tab_dir, dump_dir, papers_meta): | |
| ## Data prep | |
| import joblib | |
| # test full again - check images - check dfs !! | |
| self.clean_dirs([pdf_dir, txt_dir, img_dir, tab_dir, dump_dir]) | |
| papers = papers_meta[:self.num_papers] | |
| selected_papers = papers | |
| ids_none, papers, cites = self.fetch_papers(dump_dir, img_dir, papers, pdf_dir, tab_dir, txt_dir) | |
| self.print_fn("\n- First stage paper collection complete, papers collected: \n" + ', '.join([p['id'] for p in papers])) | |
| new_papers = papers_meta[self.num_papers : self.num_papers + len(ids_none)] | |
| # _ = self.get_freq_cited(cites) | |
| ''' | |
| filtered_idlist = [] | |
| for c in self.get_freq_cited(cites): | |
| if c in | |
| _, new_searched_papers = self.search(filtered_idlist) | |
| new_papers.extend(new_searched_papers) | |
| ''' | |
| selected_papers.extend(new_papers) | |
| _, new_papers, _ = self.fetch_papers(dump_dir, img_dir, new_papers, pdf_dir, tab_dir, txt_dir, repeat=True) | |
| self.print_fn("\n- Second stage paper collection complete, new papers collected: \n" + ', '.join([p['id'] for p in new_papers])) | |
| papers.extend(new_papers) | |
| joblib.dump(papers, dump_dir + 'papers_extracted_pdf_route.dmp') | |
| copy_tree(img_dir, dump_dir + os.path.basename(img_dir)) | |
| copy_tree(tab_dir, dump_dir + os.path.basename(tab_dir)) | |
| self.print_fn("\n- Extracting section-wise highlights.. ") | |
| papers = self.extract_highlights(papers) | |
| return papers, selected_papers, cites | |
| def get_freq_cited(self, cites_dict, k=5): | |
| cites_list = [] | |
| for k, v in cites_dict.items(): | |
| cites_list.append(k) | |
| [cites_list.append(val) for val in v] | |
| cite_freqs = {cite: cites_list.count(cite) for cite in set(cites_list)} | |
| sorted_cites = dict(sorted(cite_freqs.items(), key=lambda item: item[1], reverse=True)[:5]) | |
| return sorted_cites.keys() | |
| def fetch_papers(self, dump_dir, img_dir, papers, pdf_dir, tab_dir, txt_dir, repeat=False): | |
| import tempfile | |
| if repeat: | |
| with tempfile.TemporaryDirectory() as dirpath: | |
| self.print_fn("\n- downloading extra pdfs.. ") | |
| # full text preparation of selected papers | |
| self.download_pdfs(papers, dirpath) | |
| dirpath_pdfs = os.listdir(dirpath) | |
| for file_name in dirpath_pdfs: | |
| full_file_name = os.path.join(dirpath, file_name) | |
| if os.path.isfile(full_file_name): | |
| shutil.copy(full_file_name, pdf_dir) | |
| self.print_fn("\n- converting extra pdfs.. ") | |
| self.convert_pdfs(dirpath, txt_dir) | |
| else: | |
| self.print_fn("\n- downloading pdfs.. ") | |
| # full text preparation of selected papers | |
| self.download_pdfs(papers, pdf_dir) | |
| self.print_fn("\n- converting pdfs.. ") | |
| self.convert_pdfs(pdf_dir, txt_dir) | |
| # plugging citations to our papers object | |
| self.print_fn("\n- plugging in citation network.. ") | |
| papers, cites = self.cocitation_network(papers, txt_dir) | |
| joblib.dump(papers, dump_dir + 'papers_selected_pdf_route.dmp') | |
| from distutils.dir_util import copy_tree | |
| copy_tree(txt_dir, dump_dir + os.path.basename(txt_dir)) | |
| copy_tree(pdf_dir, dump_dir + os.path.basename(pdf_dir)) | |
| self.print_fn("\n- extracting structure.. ") | |
| papers, ids_none = self.extract_structure(papers, pdf_dir, txt_dir, img_dir, dump_dir, tab_dir) | |
| return ids_none, papers, cites | |
| def tar_route(self, pdf_dir, txt_dir, img_dir, tab_dir, papers): | |
| ## Data prep | |
| import joblib | |
| # test full again - check images - check dfs !! | |
| self.clean_dirs([pdf_dir, txt_dir, img_dir, tab_dir]) | |
| # full text preparation of selected papers | |
| self.download_sources(papers, pdf_dir) | |
| self.convert_pdfs(pdf_dir, txt_dir) | |
| # plugging citations to our papers object | |
| papers, cites = self.cocitation_network(papers, txt_dir) | |
| joblib.dump(papers, 'papers_selected_tar_route.dmp') | |
| papers = self.extract_structure(papers, pdf_dir, txt_dir, img_dir, tab_dir) | |
| joblib.dump(papers, 'papers_extracted_tar_route.dmp') | |
| return papers | |
| def build_doc(self, research_sections, papers, query=None, filename='survey.txt'): | |
| import arxiv2bib | |
| self.print_fn("\n- building bibliography entries.. ") | |
| bibentries = arxiv2bib.arxiv2bib([p['id'] for p in papers]) | |
| bibentries = [r.bibtex() for r in bibentries] | |
| self.print_fn("\n- building final survey file .. at "+ filename) | |
| file = open(filename, 'w+') | |
| if query is None: | |
| query = 'Internal(existing) research' | |
| self.survey_print_fn("#### Generated_survey:") | |
| file.write("----------------------------------------------------------------------") | |
| file.write("Title: A survey on " + query) | |
| self.survey_print_fn("") | |
| self.survey_print_fn("----------------------------------------------------------------------") | |
| self.survey_print_fn("Title: A survey on " + query) | |
| file.write("Author: Auto-Research (github.com/sidphbot/Auto-Research)") | |
| self.survey_print_fn("Author: Auto-Research (github.com/sidphbot/Auto-Research)") | |
| file.write("Dev: Auto-Research (github.com/sidphbot/Auto-Research)") | |
| self.survey_print_fn("Dev: Auto-Research (github.com/sidphbot/Auto-Research)") | |
| file.write("Disclaimer: This survey is intended to be a research starter. This Survey is Machine-Summarized, "+ | |
| "\nhence some sentences might be wrangled or grammatically incorrect. However all sentences are "+ | |
| "\nmined with proper citations. As All of the text is practically quoted texted, hence to "+ | |
| "\nimprove visibility, all the papers are duly cited in the Bibiliography section. as bibtex "+ | |
| "\nentries(only to avoid LaTex overhead). ") | |
| self.survey_print_fn("Disclaimer: This survey is intended to be a research starter. This Survey is Machine-Summarized, "+ | |
| "\nhence some sentences might be wrangled or grammatically incorrect. However all sentences are "+ | |
| "\nmined with proper citations. As All of the text is practically quoted texted, hence to "+ | |
| "\nimprove visibility, all the papers are duly cited in the Bibiliography section. as bibtex "+ | |
| "\nentries(only to avoid LaTex overhead). ") | |
| file.write("----------------------------------------------------------------------") | |
| self.survey_print_fn("----------------------------------------------------------------------") | |
| file.write("") | |
| self.survey_print_fn("") | |
| file.write('ABSTRACT') | |
| self.survey_print_fn('ABSTRACT') | |
| self.survey_print_fn("=================================================") | |
| file.write("=================================================") | |
| file.write("") | |
| self.survey_print_fn("") | |
| file.write(research_sections['abstract']) | |
| self.survey_print_fn(research_sections['abstract']) | |
| file.write("") | |
| self.survey_print_fn("") | |
| file.write('INTRODUCTION') | |
| self.survey_print_fn('INTRODUCTION') | |
| self.survey_print_fn("=================================================") | |
| file.write("=================================================") | |
| file.write("") | |
| self.survey_print_fn("") | |
| file.write(research_sections['introduction']) | |
| self.survey_print_fn(research_sections['introduction']) | |
| file.write("") | |
| self.survey_print_fn("") | |
| for k, v in research_sections.items(): | |
| if k not in ['abstract', 'introduction', 'conclusion']: | |
| file.write(k.upper()) | |
| self.survey_print_fn(k.upper()) | |
| self.survey_print_fn("=================================================") | |
| file.write("=================================================") | |
| file.write("") | |
| self.survey_print_fn("") | |
| file.write(v) | |
| self.survey_print_fn(v) | |
| file.write("") | |
| self.survey_print_fn("") | |
| file.write('CONCLUSION') | |
| self.survey_print_fn('CONCLUSION') | |
| self.survey_print_fn("=================================================") | |
| file.write("=================================================") | |
| file.write("") | |
| self.survey_print_fn("") | |
| file.write(research_sections['conclusion']) | |
| self.survey_print_fn(research_sections['conclusion']) | |
| file.write("") | |
| self.survey_print_fn("") | |
| file.write('REFERENCES') | |
| self.survey_print_fn('REFERENCES') | |
| self.survey_print_fn("=================================================") | |
| file.write("=================================================") | |
| file.write("") | |
| self.survey_print_fn("") | |
| for entry in bibentries: | |
| file.write(entry) | |
| self.survey_print_fn(entry) | |
| file.write("") | |
| self.survey_print_fn("") | |
| self.survey_print_fn("========================XXX=========================") | |
| file.write("========================XXX=========================") | |
| file.close() | |
| def build_basic_blocks(self, corpus_known_sections, corpus): | |
| research_blocks = {} | |
| for head, textarr in corpus_known_sections.items(): | |
| if 'cuda' in self.torch_device: | |
| torch.cuda.empty_cache() | |
| # self.print_fn(head.upper()) | |
| with torch.no_grad(): | |
| summtext = self.model(" ".join([l.lower() for l in textarr]), ratio=0.5) | |
| res = self.nlp(summtext) | |
| res = set([str(sent) for sent in list(res.sents)]) | |
| summtext = ''.join([line for line in res]) | |
| # pself.print_fn(summtext) | |
| research_blocks[head] = summtext | |
| return research_blocks | |
| def abstractive_summary(self, longtext): | |
| ''' | |
| faulty method | |
| input_ids = ledtokenizer(longtext, return_tensors="pt").input_ids | |
| global_attention_mask = torch.zeros_like(input_ids) | |
| # set global_attention_mask on first token | |
| global_attention_mask[:, 0] = 1 | |
| sequences = ledmodel.generate(input_ids, global_attention_mask=global_attention_mask).sequences | |
| summary = ledtokenizer.batch_decode(sequences) | |
| ''' | |
| if 'cuda' in self.torch_device: | |
| torch.cuda.empty_cache() | |
| inputs = self.ledtokenizer.prepare_seq2seq_batch(longtext, truncation=True, padding='longest', | |
| return_tensors='pt').to(self.torch_device) | |
| with torch.no_grad(): | |
| summary_ids = self.ledmodel.generate(**inputs) | |
| summary = self.ledtokenizer.batch_decode(summary_ids, skip_special_tokens=True, | |
| clean_up_tokenization_spaces=True) | |
| res = self.nlp(summary[0]) | |
| res = set([str(sent) for sent in list(res.sents)]) | |
| summtext = ''.join([line for line in res]) | |
| #self.print_fn("abstractive summary type:" + str(type(summary))) | |
| return summtext | |
| def get_abstract(self, abs_lines, corpus_known_sections, research_blocks): | |
| # abs_lines = " ".join(abs_lines) | |
| abs_lines = "" | |
| abs_lines += " ".join([l.lower() for l in corpus_known_sections['abstract']]) | |
| abs_lines += research_blocks['abstract'] | |
| # self.print_fn(abs_lines) | |
| try: | |
| return self.abstractive_summary(abs_lines) | |
| except: | |
| highlights = self.extractive_summary(abs_lines) | |
| return self.abstractive_summary(highlights) | |
| def get_corpus_lines(self, corpus): | |
| abs_lines = [] | |
| types = set() | |
| for k, v in corpus.items(): | |
| # self.print_fn(v) | |
| types.add(type(v)) | |
| abstext = k + '. ' + v.replace('\n', ' ') | |
| abstext = self.nlp(abstext) | |
| abs_lines.extend([str(sent).lower() for sent in list(abstext.sents)]) | |
| #self.print_fn("unique corpus value types:" + str(types)) | |
| # abs_lines = '\n'.join([str(sent) for sent in abs_lines.sents]) | |
| return abs_lines | |
| def get_sectioned_docs(self, papers, papers_meta): | |
| import random | |
| docs = [] | |
| for p in papers: | |
| for section in p['sections']: | |
| if len(section['highlights']) > 0: | |
| content = self.extractive_summary(''.join(section['highlights'])) | |
| docs.append(content) | |
| selected_pids = [p['id'] for p in papers] | |
| meta_abs = [] | |
| for p in papers_meta: | |
| if p['id'] not in selected_pids: | |
| meta_abs.append(self.generate_title(p['abstract'])) | |
| docs.extend(meta_abs) | |
| #self.print_fn("meta_abs num"+str(len(meta_abs))) | |
| #self.print_fn("selected_pids num"+str(len(selected_pids))) | |
| #self.print_fn("papers_meta num"+str(len(papers_meta))) | |
| #assert (len(meta_abs) + len(selected_pids) == len(papers_meta)) | |
| assert ('str' in str(type(random.sample(docs, 1)[0]))) | |
| return [doc for doc in docs if doc != ''] | |
| def cluster_lines(self, abs_lines): | |
| from sklearn.cluster import KMeans | |
| # from bertopic import BERTopic | |
| # topic_model = BERTopic(embedding_model=embedder) | |
| if 'cuda' in self.torch_device: | |
| torch.cuda.empty_cache() | |
| corpus_embeddings = self.embedder.encode(abs_lines) | |
| # Normalize the embeddings to unit length | |
| corpus_embeddings = corpus_embeddings / np.linalg.norm(corpus_embeddings, axis=1, keepdims=True) | |
| with torch.no_grad(): | |
| optimal_k = self.model.calculate_optimal_k(' '.join(abs_lines), k_max=10) | |
| # Perform kmean clustering | |
| clustering_model = KMeans(n_clusters=optimal_k, n_init=20) | |
| # clustering_model = AgglomerativeClustering(n_clusters=optimal_k, affinity='cosine', linkage='average') #, affinity='cosine', linkage='average', distance_threshold=0.4) | |
| clustering_model.fit(corpus_embeddings) | |
| cluster_assignment = clustering_model.labels_ | |
| clustered_sentences = {} | |
| dummy_count = 0 | |
| for sentence_id, cluster_id in enumerate(cluster_assignment): | |
| if cluster_id not in clustered_sentences: | |
| clustered_sentences[cluster_id] = [] | |
| ''' | |
| if dummy_count < 5: | |
| self.print_fn("abs_line: "+abs_lines[sentence_id]) | |
| self.print_fn("cluster_ID: "+str(cluster_id)) | |
| self.print_fn("embedding: "+str(corpus_embeddings[sentence_id])) | |
| dummy_count += 1 | |
| ''' | |
| clustered_sentences[cluster_id].append(abs_lines[sentence_id]) | |
| # for i, cluster in clustered_sentences.items(): | |
| # self.print_fn("Cluster ", i+1) | |
| # self.print_fn(cluster) | |
| # self.print_fn("") | |
| return self.get_clustered_sections(clustered_sentences), clustered_sentences | |
| def get_clusters(self, papers, papers_meta): | |
| from sklearn.cluster import KMeans | |
| # from bertopic import BERTopic | |
| # topic_model = BERTopic(embedding_model=embedder) | |
| if 'cuda' in self.torch_device: | |
| torch.cuda.empty_cache() | |
| abs_lines = self.get_sectioned_docs(papers, papers_meta) | |
| corpus_embeddings = self.embedder.encode(abs_lines) | |
| # Normalize the embeddings to unit length | |
| corpus_embeddings = corpus_embeddings / np.linalg.norm(corpus_embeddings, axis=1, keepdims=True) | |
| with torch.no_grad(): | |
| optimal_k = self.model.calculate_optimal_k(' '.join(abs_lines), k_max=10) | |
| # Perform kmean clustering | |
| clustering_model = KMeans(n_clusters=optimal_k, n_init=20) | |
| # clustering_model = AgglomerativeClustering(n_clusters=optimal_k, affinity='cosine', linkage='average') #, affinity='cosine', linkage='average', distance_threshold=0.4) | |
| clustering_model.fit(corpus_embeddings) | |
| cluster_assignment = clustering_model.labels_ | |
| clustered_sentences = {} | |
| dummy_count = 0 | |
| for sentence_id, cluster_id in enumerate(cluster_assignment): | |
| if cluster_id not in clustered_sentences: | |
| clustered_sentences[cluster_id] = [] | |
| ''' | |
| if dummy_count < 5: | |
| self.print_fn("abs_line: "+abs_lines[sentence_id]) | |
| self.print_fn("cluster_ID: "+str(cluster_id)) | |
| self.print_fn("embedding: "+str(corpus_embeddings[sentence_id])) | |
| dummy_count += 1 | |
| ''' | |
| clustered_sentences[cluster_id].append(abs_lines[sentence_id]) | |
| # for i, cluster in clustered_sentences.items(): | |
| # self.print_fn("Cluster ", i+1) | |
| # self.print_fn(cluster) | |
| # self.print_fn("") | |
| return self.get_clustered_sections(clustered_sentences), clustered_sentences | |
| def generate_title(self, longtext): | |
| if 'cuda' in self.torch_device: | |
| torch.cuda.empty_cache() | |
| inputs = self.title_tokenizer.prepare_seq2seq_batch(longtext, truncation=True, padding='longest', | |
| return_tensors='pt').to(self.torch_device) | |
| with torch.no_grad(): | |
| summary_ids = self.title_model.generate(**inputs) | |
| summary = self.title_tokenizer.batch_decode(summary_ids, skip_special_tokens=True, | |
| clean_up_tokenization_spaces=True) | |
| return str(summary[0]) | |
| def get_clustered_sections(self, clustered_lines): | |
| clusters_dict = {} | |
| for i, cluster in clustered_lines.items(): | |
| # self.print_fn(cluster) | |
| try: | |
| clusters_dict[self.generate_title(str(" ".join(cluster)))] = self.abstractive_summary( | |
| str(" ".join(cluster)).lower()) | |
| except: | |
| clusters_dict[self.generate_title(str(" ".join(cluster)))] = self.abstractive_summary( | |
| self.extractive_summary(str(" ".join(cluster)).lower())) | |
| return clusters_dict | |
| def get_intro(self, corpus_known_sections, research_blocks): | |
| intro_lines = "" | |
| intro_lines += str(" ".join([l.lower() for l in corpus_known_sections['introduction']])) + str( | |
| " ".join([l.lower() for l in corpus_known_sections['conclusion']])) | |
| intro_lines += research_blocks['introduction'] + research_blocks['conclusion'] | |
| try: | |
| return self.abstractive_summary(intro_lines) | |
| except: | |
| return self.abstractive_summary(self.extractive_summary(intro_lines)) | |
| def get_conclusion(self, research_sections): | |
| paper_body = "" | |
| for k, v in research_sections.items(): | |
| paper_body += v | |
| try: | |
| return self.abstractive_summary(paper_body) | |
| except: | |
| return self.abstractive_summary(self.extractive_summary(paper_body)) | |
| def build_corpus_sectionwise(self, papers): | |
| known = ['abstract', 'introduction', 'conclusion'] | |
| corpus_known_sections = {} | |
| for kh in known: | |
| khtext = [] | |
| for p in papers: | |
| for section in p['sections']: | |
| if kh in section['heading']: | |
| khtext.extend(section['highlights']) | |
| # self.print_fn(khtext) | |
| corpus_known_sections[kh] = khtext | |
| return corpus_known_sections | |
| def standardize_headings(self, papers): | |
| known = ['abstract', 'introduction', 'discussion', 'relatedwork', 'contribution', 'analysis', 'experiments', | |
| 'conclusion'] | |
| for p in papers: | |
| # self.print_fn("================================") | |
| headings = [section['heading'] for section in p['sections'] if len(section['heading'].split()) < 3] | |
| # self.print_fn("id: "+ str(p['id'])+"\nHeadings: \n"+str('\n'.join(headings))) | |
| for kh in known: | |
| for section in p['sections']: | |
| if len(section['heading'].split()) < 3: | |
| # self.print_fn(section['heading']) | |
| if kh in ''.join(filter(str.isalpha, section['heading'].replace(' ', '').lower())): | |
| # self.print_fn("orig head: "+ section['heading'] +", plain head:" + kh) | |
| section['heading'] = kh | |
| return papers | |
| def build_corpus(self, papers, papers_meta): | |
| corpus = self.build_meta_corpus(papers_meta) | |
| for p in papers: | |
| ph = [] | |
| for sid, section in enumerate(p['sections']): | |
| ph.extend(section['highlights']) | |
| for pid, ls in corpus.items(): | |
| if pid == p['id']: | |
| corpus[pid] = p['abstract'] + str(' '.join(ph)) | |
| ''' | |
| self.print_fn("================== final corpus ====================") | |
| self.print_fn('\n'.join([str("paper: "+ get_by_pid(pid, papers_meta)['title']+" \nhighlight count: " + str(len(phs))) for pid, phs in corpus.items()])) | |
| self.print_fn("======== sample point ========") | |
| p = random.choice(list(papers)) | |
| self.print_fn("paper: "+ p['title']+" \nhighlights: " + str(corpus[p['id']])) | |
| self.print_fn("======== sample meta point ========") | |
| p = random.choice(list(papers_meta)) | |
| self.print_fn("meta paper: "+ p['title']+" \nhighlights: " + str(corpus[p['id']])) | |
| ''' | |
| return corpus | |
| def get_by_pid(self, pid, papers): | |
| for p in papers: | |
| if p['id'] == pid: | |
| return p | |
| def build_meta_corpus(self, papers): | |
| meta_corpus = {} | |
| for p in papers: | |
| # pself.print_fn(p) | |
| pid = p['id'] | |
| ptext = p['title'] + ". " + p['abstract'] | |
| doc = self.nlp(ptext) | |
| phs, _, _ = self.extractive_highlights([str(sent) for sent in list(doc.sents)]) | |
| meta_corpus[pid] = str(' '.join(phs)) | |
| ''' | |
| self.print_fn("================== meta corpus ====================") | |
| self.print_fn('\n'.join([str("paper: "+ get_by_pid(pid, papers)['title']+" \nhighlight count: " + str(len(phs))) for pid, phs in meta_corpus.items()])) | |
| self.print_fn("======== sample point ========") | |
| p = random.choice(list(papers)) | |
| self.print_fn("paper: "+ p['title']+" \nhighlights: " + str(meta_corpus[p['id']])) | |
| ''' | |
| return meta_corpus | |
| def select_papers(self, papers, query, num_papers=20): | |
| import numpy as np | |
| # self.print_fn("paper sample: ") | |
| # self.print_fn(papers) | |
| meta_corpus = self.build_meta_corpus(papers) | |
| scores = [] | |
| pids = [] | |
| for id, highlights in meta_corpus.items(): | |
| score = self.text_para_similarity(query, highlights) | |
| scores.append(score) | |
| pids.append(id) | |
| # self.print_fn("corpus item: " + str(self.get_by_pid(id, papers)['title'])) | |
| idx = np.argsort(scores)[:num_papers] | |
| #for i in range(len(scores)): | |
| # self.print_fn("paper: " + str(self.get_by_pid(pids[i], papers)['title'])) | |
| # self.print_fn("score: " + str(scores[i])) | |
| # self.print_fn("argsort ids("+str(num_papers)+" papers): "+ str(idx)) | |
| idx = [pids[i] for i in idx] | |
| # self.print_fn("argsort pids("+str(num_papers)+" papers): "+ str(idx)) | |
| papers_selected = [p for p in papers if p['id'] in idx] | |
| # assert(len(papers_selected)==num_papers) | |
| self.print_fn("num papers selected: " + str(len(papers_selected))) | |
| for p in papers_selected: | |
| self.print_fn("Selected Paper: " + p['title']) | |
| #self.print_fn("constrast with natural selection: forward") | |
| #for p in papers[:4]: | |
| # self.print_fn("Selected Paper: " + p['title']) | |
| #self.print_fn("constrast with natural selection: backward") | |
| #for p in papers[-4:]: | |
| # self.print_fn("Selected Paper: " + p['title']) | |
| # arxiv search producing better relevnce | |
| return papers_selected | |
| def extractive_summary(self, text): | |
| if 'cuda' in self.torch_device: | |
| torch.cuda.empty_cache() | |
| with torch.no_grad(): | |
| res = self.model(text, ratio=0.5) | |
| res_doc = self.nlp(res) | |
| return " ".join(set([str(sent) for sent in list(res_doc.sents)])) | |
| def extractive_highlights(self, lines): | |
| # text = " ".join(lines) | |
| # text_doc = nlp(" ".join([l.lower() for l in lines])) | |
| # text = ' '.join([ str(sent) for sent in list(text_doc.sents)]) | |
| if 'cuda' in self.torch_device: | |
| torch.cuda.empty_cache() | |
| with torch.no_grad(): | |
| res = self.model(" ".join([l.lower() for l in lines]), ratio=0.5, ) | |
| res_doc = self.nlp(res) | |
| res_lines = set([str(sent) for sent in list(res_doc.sents)]) | |
| # self.print_fn("\n- ".join(res_sents)) | |
| with torch.no_grad(): | |
| keywords = self.kw_model.extract_keywords(str(" ".join([l.lower() for l in lines])), stop_words='english') | |
| keyphrases = self.kw_model.extract_keywords(str(" ".join([l.lower() for l in lines])), | |
| keyphrase_ngram_range=(4, 4), | |
| stop_words='english', use_mmr=True, diversity=0.7) | |
| return res_lines, keywords, keyphrases | |
| def extract_highlights(self, papers): | |
| for p in papers: | |
| sid = 0 | |
| p['sections'] = [] | |
| for heading, lines in p['body_text'].items(): | |
| hs, kws, kps = self.extractive_highlights(lines) | |
| p['sections'].append({ | |
| 'sid': sid, | |
| 'heading': heading, | |
| 'text': lines, | |
| 'highlights': hs, | |
| 'keywords': kws, | |
| 'keyphrases': kps, | |
| }) | |
| sid += 1 | |
| return papers | |
| def extract_structure(self, papers, pdf_dir, txt_dir, img_dir, dump_dir, tab_dir, tables=False): | |
| self.print_fn("\n- extracting sections.. ") | |
| papers, ids_none = self.extract_parts(papers, txt_dir, dump_dir) | |
| self.print_fn("\n- extracting images.. for future correlation use-cases ") | |
| papers = self.extract_images(papers, pdf_dir, img_dir) | |
| if tables: | |
| self.print_fn("\n- extracting tables.. for future correlation use-cases ") | |
| papers = self.extract_tables(papers, pdf_dir, tab_dir) | |
| return papers, ids_none | |
| def extract_parts(self, papers, txt_dir, dump_dir): | |
| headings_all = {} | |
| # refined = [] | |
| # model = build_summarizer() | |
| #for file in glob.glob(txt_dir + '/*.txt'): | |
| for p in papers: | |
| file = txt_dir + '/'+ p['id'] +'.txt' | |
| refined, headings_extracted = self.extract_headings(file) | |
| sections = self.extract_sections(headings_extracted, refined) | |
| # highlights = {k: extract_highlights(model,v) for k, v in sections.items()} | |
| #p = self.get_by_file(file, papers) | |
| #if len(headings_extracted) > 3: | |
| p['body_text'] = sections | |
| # p['body_highlights'] = highlights | |
| headings_all[p['id']] = headings_extracted | |
| ids_none = {i: h for i, h in headings_all.items() if len(h) < 3} | |
| ''' | |
| for f, h in headings_all.items(): | |
| if len(h) < 4: | |
| self.print_fn("=================headings almost undetected================") | |
| self.print_fn(f) | |
| self.print_fn(h) | |
| ''' | |
| # from pprint import pprint | |
| # pself.print_fn({f: len(h) for f,h in headings_all.items()}) | |
| papers_none = [p for p in papers if p['id'] in ids_none] | |
| for p in papers_none: | |
| os.remove(txt_dir + '/'+ p['id'] + '.txt') | |
| papers.remove(p) | |
| return papers, ids_none | |
| def check_para(self, df): | |
| size = 0 | |
| for col in df.columns: | |
| size += df[col].apply(lambda x: len(str(x))).median() | |
| return size / len(df.columns) > 25 | |
| def scan_blocks(self, lines): | |
| lines_mod = [line.strip().replace('\n', '') for line in lines if len(line.strip().replace('\n', '')) > 3] | |
| for i in range(len(lines_mod)): | |
| yield lines_mod[i:i + 3] | |
| def extract_sections(self, headings, lines, min_part_length=2): | |
| sections = {} | |
| self.check_list_elems_in_list(headings, lines) | |
| head_len = len(headings) | |
| for i in range(len(headings) - 1): | |
| start = headings[i] | |
| end = headings[i + 1] | |
| section = self.get_section(start, end, lines) | |
| # self.print_fn(start + " : "+ str(len(section)) +" lines") | |
| ''' | |
| if i > 0: | |
| old = headings[i-1] | |
| if len(section) < min_part_length + 1: | |
| sections[old].extend(start) | |
| sections[old].extend(section) | |
| else: | |
| sections[start] = section | |
| else: | |
| sections[start] = section | |
| ''' | |
| sections[start] = section | |
| return {k: v for k, v in sections.items()} | |
| def is_rubbish(self, s, rubbish_tolerance=0.2, min_char_len=4): | |
| # numbers = sum(c.isdigit() for c in s) | |
| letters = sum(c.isalpha() for c in s) | |
| spaces = sum(c.isspace() for c in s) | |
| # others = len(s) - numbers - letters - spaces | |
| if len(s) == 0: | |
| return False | |
| if ((len(s) - (letters + spaces)) / len(s) >= rubbish_tolerance) or self.alpha_length(s) < min_char_len: | |
| return True | |
| else: | |
| return False | |
| def get_section(self, first, last, lines): | |
| try: | |
| assert (first in lines) | |
| assert (last in lines) | |
| # start = lines.index( first ) + len( first ) | |
| # end = lines.index( last, start ) | |
| start = [i for i in range(len(lines)) if first is lines[i]][0] | |
| end = [i for i in range(len(lines)) if last is lines[i]][0] | |
| section_lines = lines[start + 1:end] | |
| # self.print_fn("heading: " + str(first)) | |
| # self.print_fn("section_lines: "+ str(section_lines)) | |
| # self.print_fn(section_lines) | |
| return section_lines | |
| except ValueError: | |
| self.print_fn("value error :") | |
| self.print_fn("first heading :" + str(first) + ", second heading :" + str(last)) | |
| self.print_fn("first index :" + str(start) + ", second index :" + str(end)) | |
| return "" | |
| def check_list_elems_in_list(self, headings, lines): | |
| import numpy as np | |
| # [self.print_fn(head) for head in headings if head not in lines ] | |
| return np.all([True if head in lines else False for head in headings]) | |
| def check_first_char_upper(self, text): | |
| for c in text: | |
| if c.isspace(): | |
| continue | |
| elif c.isalpha(): | |
| return c.isupper() | |
| def extract_headings(self, txt_file): | |
| import re | |
| fulltext = self.read_paper(txt_file) | |
| lines = self.clean_lines(fulltext) | |
| refined, headings = self.scan_text(lines) | |
| assert (self.check_list_elems_in_list(headings, refined)) | |
| headings = self.check_duplicates(headings) | |
| # self.print_fn('===========================================') | |
| # self.print_fn(txt_file +": first scan: \n"+str(len(headings))+" headings") | |
| # self.print_fn('\n'.join(headings)) | |
| # scan_failed - rescan with first match for abstract hook | |
| if len(headings) == 0: | |
| # self.print_fn('===================') | |
| # self.print_fn("run 1 failed") | |
| abs_cans = [line for line in lines if 'abstract' in re.sub("\s+", "", line.strip().lower())] | |
| if len(abs_cans) != 0: | |
| abs_head = abs_cans[0] | |
| refined, headings = self.scan_text(lines, abs_head=abs_head) | |
| self.check_list_elems_in_list(headings, refined) | |
| headings = self.check_duplicates(headings) | |
| # self.print_fn('===================') | |
| # self.print_fn(txt_file +": second scan: \n"+str(len(headings))+" headings") | |
| # if len(headings) == 0: | |
| # self.print_fn("heading scan failed completely") | |
| return refined, headings | |
| def check_duplicates(self, my_list): | |
| my_finallist = [] | |
| dups = [s for s in my_list if my_list.count(s) > 1] | |
| if len(dups) > 0: | |
| [my_finallist.append(n) for n in my_list if n not in my_finallist] | |
| # self.print_fn("original: "+str(len(my_list))+" new: "+str(len(my_finallist))) | |
| return my_finallist | |
| def clean_lines(self, text): | |
| import numpy as np | |
| import re | |
| # doc = nlp(text) | |
| # lines = [str(sent) for sent in doc.sents] | |
| lines = text.replace('\r', '').split('\n') | |
| lines = [line for line in lines if not self.is_rubbish(line)] | |
| lines = [line for line in lines if | |
| re.match("^[a-zA-Z1-9\.\[\]\(\):\-,\"\"\s]*$", line) and not 'Figure' in line and not 'Table' in line] | |
| lengths_cleaned = [self.alpha_length(line) for line in lines] | |
| mean_length_cleaned = np.median(lengths_cleaned) | |
| lines_standardized = [] | |
| for line in lines: | |
| if len(line) >= (1.8 * mean_length_cleaned): | |
| first_half = line[0:len(line) // 2] | |
| second_half = line[len(line) // 2 if len(line) % 2 == 0 else ((len(line) // 2) + 1):] | |
| lines_standardized.append(first_half) | |
| lines_standardized.append(second_half) | |
| else: | |
| lines_standardized.append(line) | |
| return lines | |
| def scan_text(self, lines, abs_head=None): | |
| import re | |
| # self.print_fn('\n'.join(lines)) | |
| record = False | |
| headings = [] | |
| refined = [] | |
| for i in range(1, len(lines) - 4): | |
| line = lines[i] | |
| line = line.replace('\n', '').strip() | |
| if 'abstract' in re.sub("\s+", "", line.strip().lower()) and len(line) - len('abstract') < 5 or ( | |
| abs_head is not None and abs_head in line): | |
| record = True | |
| headings.append(line) | |
| refined.append(line) | |
| if 'references' in re.sub("\s+", "", line.strip().lower()) and len(line) - len('references') < 5: | |
| headings.append(line) | |
| refined.append(line) | |
| break | |
| elif 'bibliography' in re.sub("\s+", "", line.strip().lower()) and len(line) - len('bibliography') < 5: | |
| headings.append(line) | |
| refined.append(line) | |
| break | |
| refined, headings = self.scanline(record, headings, refined, i, lines) | |
| # self.print_fn('=========in scan_text loop i : '+str(i)+' heading count : '+str(len(headings))+' =========') | |
| return refined, headings | |
| def scanline(self, record, headings, refined, id, lines): | |
| import numpy as np | |
| import re | |
| line = lines[id] | |
| if not len(line) == 0: | |
| # self.print_fn("in scanline") | |
| # self.print_fn(line) | |
| if record: | |
| refined.append(line) | |
| if len(lines[id - 1]) == 0 or len(lines[id + 1]) == 0 or re.match( | |
| "^[1-9XVIABCD]{0,4}(\.{0,1}[1-9XVIABCD]{0,4}){0, 3}\s{0,2}[A-Z][a-zA-Z\:\-\s]*$", | |
| line) and self.char_length(line) > 7: | |
| # self.print_fn("candidate") | |
| # self.print_fn(line) | |
| if np.mean([len(s) for s in lines[id + 2:id + 6]]) > 40 and self.check_first_char_upper( | |
| line) and re.match("^[a-zA-Z1-9\.\:\-\s]*$", line) and len(line.split()) < 10: | |
| # if len(line) < 20 and np.mean([len(s) for s in lines[i+1:i+5]]) > 30 : | |
| headings.append(line) | |
| assert (line in refined) | |
| # self.print_fn("selected") | |
| # self.print_fn(line) | |
| else: | |
| known_headings = ['introduction', 'conclusion', 'abstract', 'references', 'bibliography'] | |
| missing = [h for h in known_headings if not np.any([True for head in headings if h in head])] | |
| # for h in missing: | |
| head = [line for h in missing if h in re.sub("\s+", "", line.strip().lower())] | |
| # head = [line for known] | |
| if len(head) > 0: | |
| headings.append(head[0]) | |
| assert (head[0] in refined) | |
| return refined, headings | |
| def char_length(self, s): | |
| # numbers = sum(c.isdigit() for c in s) | |
| letters = sum(c.isalpha() for c in s) | |
| # spaces = sum(c.isspace() for c in s) | |
| # others = len(s) - numbers - letters - spaces | |
| return letters | |
| def get_by_file(self, file, papers): | |
| import os | |
| pid = os.path.basename(file) | |
| pid = pid.replace('.txt', '').replace('.pdf', '') | |
| for p in papers: | |
| if p['id'] == pid: | |
| return p | |
| self.print_fn("\n- paper not found by file, \nfile: "+file+"\nall papers: "+', '.join([p['id'] for p in papers])) | |
| def alpha_length(self, s): | |
| # numbers = sum(c.isdigit() for c in s) | |
| letters = sum(c.isalpha() for c in s) | |
| spaces = sum(c.isspace() for c in s) | |
| # others = len(s) - numbers - letters - spaces | |
| return letters + spaces | |
| def check_append(self, baselist, addstr): | |
| check = False | |
| for e in baselist: | |
| if addstr in e: | |
| check = True | |
| if not check: | |
| baselist.append(addstr) | |
| return baselist | |
| def extract_images(self, papers, pdf_dir, img_dir): | |
| import fitz | |
| # self.print_fn("in images") | |
| for p in papers: | |
| file = pdf_dir + p['id'] + ".pdf" | |
| print(Path(pdf_dir).resolve().glob('*')) | |
| pdf_file = fitz.open(file) | |
| images = [] | |
| for page_index in range(len(pdf_file)): | |
| page = pdf_file[page_index] | |
| images.extend(page.getImageList()) | |
| images_files = [self.save_image(pdf_file.extractImage(img[0]), i, p['id'], img_dir) for i, img in | |
| enumerate(set(images)) if img[0]] | |
| # self.print_fn(len(images_per_paper)) | |
| p['images'] = images_files | |
| # self.print_fn(len(p.keys())) | |
| # self.print_fn(papers[0].keys()) | |
| return papers | |
| def extract_images_from_file(self, pdf_file_name, img_dir): | |
| import fitz | |
| pdf_file = fitz.open(pdf_file_name) | |
| images = [] | |
| for page_index in range(len(pdf_file)): | |
| page = pdf_file[page_index] | |
| images.extend(page.getImageList()) | |
| images_files = [self.save_image(pdf_file.extractImage(img[0]), i, pdf_file_name.replace('.pdf', ''), img_dir) for i, img in | |
| enumerate(set(images)) if img[0]] | |
| return images_files | |
| def save_image(self, base_image, img_index, pid, img_dir): | |
| from PIL import Image | |
| import io | |
| image_bytes = base_image["image"] | |
| # get the image extension | |
| image_ext = base_image["ext"] | |
| # load it to PIL | |
| image = Image.open(io.BytesIO(image_bytes)) | |
| # save it to local disk | |
| fname = img_dir + "/" + str(pid) + "_" + str(img_index + 1) + "." + image_ext | |
| image.save(open(f"{fname}", "wb")) | |
| # self.print_fn(fname) | |
| return fname | |
| def save_tables(self, dfs, pid, tab_dir): | |
| # todo | |
| dfs = [df for df in dfs if not self.check_para(df)] | |
| files = [] | |
| for df in dfs: | |
| filename = tab_dir + "/" + str(pid) + ".csv" | |
| files.append(filename) | |
| df.to_csv(filename, index=False) | |
| return files | |
| def extract_tables(self, papers, pdf_dir, tab_dir): | |
| import tabula | |
| check = True | |
| # for file in glob.glob(pdf_dir+'/*.pdf'): | |
| for p in papers: | |
| dfs = tabula.read_pdf(pdf_dir + p['id'] + ".pdf", pages='all', multiple_tables=True, silent=True) | |
| p['tables'] = self.save_tables(dfs, p['id'], tab_dir) | |
| # self.print_fn(papers[0].keys()) | |
| return papers | |
| def extract_tables_from_file(self, pdf_file_name, tab_dir): | |
| import tabula | |
| check = True | |
| # for file in glob.glob(pdf_dir+'/*.pdf'): | |
| dfs = tabula.read_pdf(pdf_file_name, pages='all', multiple_tables=True, silent=True) | |
| return self.save_tables(dfs, pdf_file_name.replace('.pdf', ''), tab_dir) | |
| def search(self, query_text=None, id_list=None, max_search=100): | |
| import arxiv | |
| from urllib.parse import urlparse | |
| if query_text: | |
| search = arxiv.Search( | |
| query=query_text, | |
| max_results=max_search, | |
| sort_by=arxiv.SortCriterion.Relevance | |
| ) | |
| else: | |
| id_list = [id for id in id_list if '.' in id] | |
| search = arxiv.Search( | |
| id_list=id_list | |
| ) | |
| results = [result for result in search.get()] | |
| searched_papers = [] | |
| discarded_ids = [] | |
| for result in results: | |
| id = urlparse(result.entry_id).path.split('/')[-1].split('v')[0] | |
| if '.' in id: | |
| paper = { | |
| 'id': id, | |
| 'title': result.title, | |
| 'comments': result.comment if result.journal_ref else "None", | |
| 'journal-ref': result.journal_ref if result.journal_ref else "None", | |
| 'doi': str(result.doi), | |
| 'primary_category': result.primary_category, | |
| 'categories': result.categories, | |
| 'license': None, | |
| 'abstract': result.summary, | |
| 'published': result.published, | |
| 'pdf_url': result.pdf_url, | |
| 'links': [str(l) for l in result.links], | |
| 'update_date': result.updated, | |
| 'authors': [str(a.name) for a in result.authors], | |
| } | |
| searched_papers.append(paper) | |
| else: | |
| discarded_ids.append(urlparse(result.entry_id).path.split('/')[-1].split('v')[0]) | |
| self.print_fn("\n- Papers discarded due to id error [arxiv api bug: #74] :\n" + str(discarded_ids)) | |
| return results, searched_papers | |
| def download_pdfs(self, papers, pdf_dir): | |
| import arxiv | |
| from urllib.parse import urlparse | |
| ids = [p['id'] for p in papers] | |
| # asert(False) | |
| papers_filtered = arxiv.Search(id_list=ids).get() | |
| for p in papers_filtered: | |
| p_id = str(urlparse(p.entry_id).path.split('/')[-1]).split('v')[0] | |
| download_file = pdf_dir + "/" + p_id + ".pdf" | |
| p.download_pdf(filename=download_file) | |
| def download_sources(self, papers, src_dir): | |
| import arxiv | |
| from urllib.parse import urlparse | |
| ids = [p['id'] for p in papers] | |
| # asert(False) | |
| papers_filtered = arxiv.Search(id_list=ids).get() | |
| for p in papers_filtered: | |
| p_id = str(urlparse(p.entry_id).path.split('/')[-1]).split('v')[0] | |
| download_file = src_dir + "/" + p_id + ".tar.gz" | |
| p.download_source(filename=download_file) | |
| def convert_pdfs(self, pdf_dir, txt_dir): | |
| import glob, shutil | |
| import multiprocessing | |
| # import arxiv_public_data | |
| convert_directory_parallel(pdf_dir, multiprocessing.cpu_count()) | |
| for file in glob.glob(pdf_dir + '/*.txt'): | |
| shutil.move(file, txt_dir) | |
| def read_paper(self, path): | |
| f = open(path, 'r', encoding="utf-8") | |
| text = str(f.read()) | |
| f.close() | |
| return text | |
| def cocitation_network(self, papers, txt_dir): | |
| import multiprocessing | |
| cites = internal_citations.citation_list_parallel(N=multiprocessing.cpu_count(), directory=txt_dir) | |
| for p in papers: | |
| p['cites'] = cites[p['id']] | |
| return papers, cites | |
| def lookup_author(self, author_query): | |
| from scholarly import scholarly | |
| import operator | |
| # Retrieve the author's data, fill-in, and print | |
| self.print_fn("Searching Author: " + author_query) | |
| search_result = next(scholarly.search_author(author_query), None) | |
| if search_result is not None: | |
| author = scholarly.fill(search_result) | |
| author_stats = { | |
| 'name': author_query, | |
| 'affiliation': author['affiliation'] if author['affiliation'] else None, | |
| 'citedby': author['citedby'] if 'citedby' in author.keys() else 0, | |
| 'most_cited_year': max(author['cites_per_year'].items(), key=operator.itemgetter(1))[0] if len( | |
| author['cites_per_year']) > 0 else None, | |
| 'coauthors': [c['name'] for c in author['coauthors']], | |
| 'hindex': author['hindex'], | |
| 'impact': author['i10index'], | |
| 'interests': author['interests'], | |
| 'publications': [{'title': p['bib']['title'], 'citations': p['num_citations']} for p in | |
| author['publications']], | |
| 'url_picture': author['url_picture'], | |
| } | |
| else: | |
| self.print_fn("author not found") | |
| author_stats = { | |
| 'name': author_query, | |
| 'affiliation': "", | |
| 'citedby': 0, | |
| 'most_cited_year': None, | |
| 'coauthors': [], | |
| 'hindex': 0, | |
| 'impact': 0, | |
| 'interests': [], | |
| 'publications': [], | |
| 'url_picture': "", | |
| } | |
| # pself.print_fn(author_stats) | |
| return author_stats | |
| def author_stats(self, papers): | |
| all_authors = [] | |
| for p in papers: | |
| paper_authors = [a for a in p['authors']] | |
| all_authors.extend(paper_authors) | |
| searched_authors = [self.lookup_author(a) for a in set(all_authors)] | |
| return searched_authors | |
| def text_similarity(self, text1, text2): | |
| doc1 = self.similarity_nlp(text1) | |
| doc2 = self.similarity_nlp(text2) | |
| return doc1.similarity(doc2) | |
| def text_para_similarity(self, text, lines): | |
| doc1 = self.similarity_nlp(text) | |
| doc2 = self.similarity_nlp(" ".join(lines)) | |
| return doc1.similarity(doc2) | |
| def para_para_similarity(self, lines1, lines2): | |
| doc1 = self.similarity_nlp(" ".join(lines1)) | |
| doc2 = self.similarity_nlp(" ".join(lines2)) | |
| return doc1.similarity(doc2) | |
| def text_image_similarity(self, text, image): | |
| pass | |
| def ask(self, corpus, question): | |
| text = " ".join(corpus) | |
| import torch | |
| inputs = self.qatokenizer(question, text, return_tensors='pt') | |
| start_positions = torch.tensor([1]) | |
| end_positions = torch.tensor([3]) | |
| outputs = self.qamodel(**inputs, start_positions=start_positions, end_positions=end_positions) | |
| self.print_fn("context: " + text) | |
| self.print_fn("question: " + question) | |
| self.print_fn("outputs: " + outputs) | |
| return outputs | |
| def zip_outputs(self, dump_dir, zip_name): | |
| import zipfile | |
| def zipdir(path, ziph): | |
| # ziph is zipfile handle | |
| for root, dirs, files in os.walk(path): | |
| for file in files: | |
| ziph.write(os.path.join(root, file), | |
| os.path.relpath(os.path.join(root, file), | |
| os.path.join(path, '../..'))) | |
| zipf = zipfile.ZipFile(zip_name, 'w', zipfile.ZIP_DEFLATED) | |
| zipdir(dump_dir, zipf) | |
| def survey(self, query=None, id_list=None, max_search=None, num_papers=None, debug=False, weigh_authors=False, | |
| pdf_dir=None, txt_dir=None, img_dir=None, tab_dir=None, dump_dir=None): | |
| import joblib | |
| import os, shutil | |
| dirs = self.define_structure(pdf_dir=pdf_dir, txt_dir=txt_dir, img_dir=img_dir, tab_dir=tab_dir, dump_dir=dump_dir) | |
| [survey_pdf_dir, survey_txt_dir, survey_img_dir, survey_tab_dir, survey_dump_dir] = dirs | |
| if not max_search: | |
| max_search = self.DEFAULTS['max_search'] | |
| if not num_papers: | |
| num_papers = self.DEFAULTS['num_papers'] | |
| if (query is None) and (id_list is None): | |
| raise ValueError('please provide a base to survey on: list of arxiv IDs or a few research keywords') | |
| # arxiv api relevance search and data preparation | |
| self.print_fn("\n- searching arXiv for top 100 papers.. ") | |
| results, searched_papers = self.search(query, id_list, max_search=max_search) | |
| joblib.dump(searched_papers, survey_dump_dir + 'papers_metadata.dmp') | |
| self.print_fn("\n- found " + str(len(searched_papers)) + " papers") | |
| # paper selection by scibert vector embedding relevance scores | |
| # papers_selected = select_papers(searched_papers, query, num_papers=num_papers) | |
| papers_highlighted, papers_selected, cites = self.pdf_route(survey_pdf_dir, survey_txt_dir, survey_img_dir, survey_tab_dir, survey_dump_dir, | |
| searched_papers) | |
| if weigh_authors: | |
| authors = self.author_stats(papers_highlighted) | |
| joblib.dump(papers_highlighted, survey_dump_dir + 'papers_highlighted.dmp') | |
| self.print_fn("\n- Standardizing known section headings per paper.. ") | |
| papers_standardized = self.standardize_headings(papers_highlighted) | |
| joblib.dump(papers_standardized, survey_dump_dir + 'papers_standardized.dmp') | |
| self.print_fn("\n- Building paper-wise corpus.. ") | |
| corpus = self.build_corpus(papers_highlighted, searched_papers) | |
| joblib.dump(corpus, survey_dump_dir + 'corpus.dmp') | |
| self.print_fn("\n- Building section-wise corpus.. ") | |
| corpus_sectionwise = self.build_corpus_sectionwise(papers_standardized) | |
| joblib.dump(corpus_sectionwise, survey_dump_dir + 'corpus_sectionwise.dmp') | |
| self.print_fn("\n- Building basic research highlights.. ") | |
| research_blocks = self.build_basic_blocks(corpus_sectionwise, corpus) | |
| joblib.dump(research_blocks, survey_dump_dir + 'research_blocks.dmp') | |
| self.print_fn("\n- Reducing corpus to lines.. ") | |
| corpus_lines = self.get_corpus_lines(corpus) | |
| joblib.dump(corpus_lines, survey_dump_dir + 'corpus_lines.dmp') | |
| # temp | |
| # searched_papers = joblib.load(dump_dir + 'papers_metadata.dmp') | |
| ''' | |
| papers_highlighted = joblib.load(dump_dir + 'papers_highlighted.dmp') | |
| corpus = joblib.load(dump_dir + 'corpus.dmp') | |
| papers_standardized = joblib.load(dump_dir + 'papers_standardized.dmp') | |
| corpus_sectionwise = joblib.load(dump_dir + 'corpus_sectionwise.dmp') | |
| research_blocks = joblib.load(dump_dir + 'research_blocks.dmp') | |
| corpus_lines = joblib.load(dump_dir + 'corpus_lines.dmp') | |
| ''' | |
| ''' | |
| self.print_fn("papers_highlighted types:"+ str(np.unique([str(type(p['sections'][0]['highlights'])) for p in papers_highlighted]))) | |
| self.print_fn("papers_highlighted example:") | |
| self.print_fn(random.sample(list(papers_highlighted), 1)[0]['sections'][0]['highlights']) | |
| self.print_fn("corpus types:"+ str(np.unique([str(type(txt)) for k,txt in corpus.items()]))) | |
| self.print_fn("corpus example:") | |
| self.print_fn(random.sample(list(corpus.items()), 1)[0]) | |
| self.print_fn("corpus_lines types:"+ str(np.unique([str(type(txt)) for txt in corpus_lines]))) | |
| self.print_fn("corpus_lines example:") | |
| self.print_fn(random.sample(list(corpus_lines), 1)[0]) | |
| self.print_fn("corpus_sectionwise types:"+ str(np.unique([str(type(txt)) for k,txt in corpus_sectionwise.items()]))) | |
| self.print_fn("corpus_sectionwise example:") | |
| self.print_fn(random.sample(list(corpus_sectionwise.items()), 1)[0]) | |
| self.print_fn("research_blocks types:"+ str(np.unique([str(type(txt)) for k,txt in research_blocks.items()]))) | |
| self.print_fn("research_blocks example:") | |
| self.print_fn(random.sample(list(research_blocks.items()), 1)[0]) | |
| ''' | |
| # self.print_fn("corpus types:"+ str(np.unique([type(txt) for k,txt in corpus.items()]))) | |
| self.print_fn("\n- Building abstract.. ") | |
| abstract_block = self.get_abstract(corpus_lines, corpus_sectionwise, research_blocks) | |
| joblib.dump(abstract_block, survey_dump_dir + 'abstract_block.dmp') | |
| ''' | |
| self.print_fn("abstract_block type:"+ str(type(abstract_block))) | |
| self.print_fn("abstract_block:") | |
| self.print_fn(abstract_block) | |
| ''' | |
| self.print_fn("\n- Building introduction.. ") | |
| intro_block = self.get_intro(corpus_sectionwise, research_blocks) | |
| joblib.dump(intro_block, survey_dump_dir + 'intro_block.dmp') | |
| ''' | |
| self.print_fn("intro_block type:"+ str(type(intro_block))) | |
| self.print_fn("intro_block:") | |
| self.print_fn(intro_block) | |
| ''' | |
| self.print_fn("\n- Building custom sections.. ") | |
| clustered_sections, clustered_sentences = self.get_clusters(papers_standardized, searched_papers) | |
| joblib.dump(clustered_sections, survey_dump_dir + 'clustered_sections.dmp') | |
| joblib.dump(clustered_sentences, survey_dump_dir + 'clustered_sentences.dmp') | |
| ''' | |
| self.print_fn("clusters extracted") | |
| self.print_fn("clustered_sentences types:"+ str(np.unique([str(type(txt)) for k,txt in clustered_sentences.items()]))) | |
| self.print_fn("clustered_sentences example:") | |
| self.print_fn(random.sample(list(clustered_sections.items()), 1)[0]) | |
| self.print_fn("clustered_sections types:"+ str(np.unique([str(type(txt)) for k,txt in clustered_sections.items()]))) | |
| self.print_fn("clustered_sections example:") | |
| self.print_fn(random.sample(list(clustered_sections.items()), 1)[0]) | |
| ''' | |
| clustered_sections['abstract'] = abstract_block | |
| clustered_sections['introduction'] = intro_block | |
| joblib.dump(clustered_sections, survey_dump_dir + 'research_sections.dmp') | |
| self.print_fn("\n- Building conclusion.. ") | |
| conclusion_block = self.get_conclusion(clustered_sections) | |
| joblib.dump(conclusion_block, survey_dump_dir + 'conclusion_block.dmp') | |
| clustered_sections['conclusion'] = conclusion_block | |
| ''' | |
| self.print_fn("conclusion_block type:"+ str(type(conclusion_block))) | |
| self.print_fn("conclusion_block:") | |
| self.print_fn(conclusion_block) | |
| ''' | |
| if query is None: | |
| query = self.generate_title(' '.join([v for v in clustered_sections.values()])) | |
| survey_file = 'A_Survey_on_' + query.replace(' ', '_') + '.txt' | |
| survey_file = Path(survey_dump_dir).resolve() / survey_file | |
| self.build_doc(clustered_sections, papers_standardized, query=query, filename=str(survey_file)) | |
| self.survey_print_fn("\n-citation-network: ") | |
| self.survey_print_fn(cites) | |
| for d in [survey_pdf_dir, survey_txt_dir, survey_img_dir, survey_tab_dir]: | |
| shutil.copytree(str(d), survey_dump_dir + f'/{Path(d).stem}/') | |
| assert (os.path.exists(survey_file)) | |
| zip_name = 'arxiv_dumps_'+query.replace(' ', '_')+'.zip' | |
| zip_name = Path(survey_dump_dir).parent.resolve() / zip_name | |
| self.zip_outputs(survey_dump_dir, str(zip_name)) | |
| self.print_fn("\n- Survey complete.. \nSurvey file path :" + str(survey_file) + | |
| "\nAll outputs zip path :" + str(zip_name)) | |
| return str(zip_name.resolve()), str(survey_file.resolve()) | |