pubmad.pipeline

  0import os
  1import networkx as nx
  2import networkx.algorithms.community as nx_comm
  3from pubmad.utils import download_articles, extract_entities, extract_naive_relations, extract_biobert_relations
  4from pubmad.types import Article, Entity
  5from typing import List, Tuple
  6from datetime import datetime
  7import time
  8from tqdm import tqdm
  9
 10def get_communities(G: nx.Graph, weight_label: str = 'weight', seed: int = 42) -> List[List[str]]:
 11    '''
 12    Returns the communities of the graph.
 13    \nArgs:
 14        \nG (nx.Graph): The graph to be used.
 15        \nweight_label (str): The weight label to be used. Defaults to 'weight'.
 16        \nseed (int): The seed to be used. Defaults to 42.
 17    \nReturns:
 18        \nList[List[str]]: A list of communities in the graph, each community is defined by a list of the mesh_ids of the nodes.
 19    '''
 20    comm = nx_comm.louvain_communities(G, weight_label, seed=seed)
 21    comm = [list(c) for c in comm if len(list(c)) > 1]
 22    return comm
 23
 24def get_graph(query: str, max_publications: int = 10, start_year: int = 1800, end_year: int = datetime.now().year, use_biobert: bool = True, save_graph: bool = True, G: nx.Graph = None, clear_cache: bool = False, callback_fn = lambda x: None, source: str = 'abstract') -> nx.Graph:
 25    '''
 26    Returns a networkx graph containing relationships between genes and diseases and genes and drugs
 27    \nArgs:
 28        \nquery (str): The query to be used to search for articles in PubMed.
 29        \nmax_publications (int): The maximum number of publications to be used. Defaults to 10.
 30        \nstart_year (int): The start year to be used. Defaults to 1800.
 31        \nend_year (int): The end year to be used. Defaults to the current year.
 32        \nuse_biobert (bool): Whether to use BioBERT to extract relations. Defaults to True. Otherwise it extracts relations using the naive approach of co-occurrence.
 33        \nsave_graph (bool): Whether to save the graph to disk. Defaults to True. It uses the .graphml extension which can be imported into Cytoscape.
 34        \nG (nx.Graph): The graph to be used. Defaults to None. If None, a new graph is created.
 35        \nclear_cache (bool): Whether to clear the cache. Defaults to False.
 36        \ncallback_fn (lambda): A callback function which receives the partial graph as an argument. Defaults to None.
 37        \nsource (str): The source to be used. Defaults to 'abstract'.
 38    \nReturns:
 39        \nnx.Graph: A networkx graph containing relationships between genes and diseases and genes and drugs.
 40    '''
 41    if G is None:
 42        G = nx.Graph()
 43        G.clear()
 44
 45    # Download the articles using pymed
 46    articles: List[Article] = download_articles(title=query, start_year=start_year, end_year=end_year, max_results = max_publications)
 47
 48    bern_calls_counter = 1
 49    i = 0 # current article
 50    N = len(articles)
 51    start = time.time()
 52
 53    for i in tqdm(range(0,len(articles))):
 54        try:
 55            article = articles[i]
 56            if use_biobert == False:
 57                entities: List[Entity] = extract_entities(article, source)
 58
 59                relations: List[Tuple[Entity, Entity, float]] = extract_naive_relations(entities)
 60            else:
 61                entities, relations, used_cache = extract_biobert_relations(article, source, clear_cache)
 62                if not used_cache:
 63                    bern_calls_counter += 1
 64                    elapsed = time.time() - start
 65                    if elapsed < 0.34:
 66                        time.sleep(0.35 - elapsed)
 67                start = time.time()
 68            
 69            # Add the entities to the graph as nodes
 70            for entity in entities:
 71                # Search if there is already a node with the same mesh_id
 72                if entity.mesh_id[0] not in G.nodes:
 73                    G.add_node(entity.mesh_id[0], mention=entity.mention, type=entity.type, pmid=entity.pmid)
 74                else:
 75                    # If there is already a node with the same mesh_id, add the mention to the node
 76                    #G.nodes[entity.mesh_id[0]]['mention'] += ' ' + entity.mention
 77                    G.nodes[entity.mesh_id[0]]['pmid'] += ',' + entity.pmid
 78
 79            # Add the relations to the graph as edges
 80            for src, dst, weight in relations:
 81                G.add_edge(src.mesh_id[0], dst.mesh_id[0], weight=weight)
 82            callback_fn(G)
 83            i += 1
 84        except Exception as e:
 85            print(f'Error: {e}')
 86            print('Skipping article')
 87            i += 1
 88
 89    # Save the graph in cytoscape format
 90    if save_graph == True:
 91        #create output directory if not exist
 92        if not os.path.exists("output"):
 93            # Create a new directory because it does not exist 
 94            os.makedirs("output")
 95            print("The new directory is created!")
 96 
 97        file_name = f"output/{query}_{start_year}_{end_year}_{max_publications}_{source}_{'BioBert' if use_biobert else 'Naive'}.graphml"
 98        
 99        
100        nx.write_graphml(G, file_name)
101        print("Graph saved in:", file_name)
102    
103    return G
104
105def add_nodes(graph: nx.Graph, query: str, max_publications: int = 10, start_year: int = 1800, end_year: int = datetime.now().year, use_biobert: bool = True, save_graph: bool = True) -> nx.Graph:
106    '''
107    Adds nodes to the graph.
108    \nArgs:
109        \ngraph (nx.Graph): The graph to be used.
110        \nquery (str): The query to be used to search for articles in PubMed.
111        \nmax_publications (int): The maximum number of publications to be used. Defaults to 10.
112        \nstart_year (int): The start year to be used. Defaults to 1800.
113        \nend_year (int): The end year to be used. Defaults to the current year.
114        \nuse_biobert (bool): Whether to use BioBERT to extract relations. Defaults to True. Otherwise it extracts relations using the naive approach of co-occurrence.
115        \nsave_graph (bool): Whether to save the graph to disk. Defaults to True. It uses the .graphml extension which can be imported into Cytoscape.
116    \nReturns:
117        \nnx.Graph: A networkx graph containing relationships between genes and diseases and genes and drugs.
118    '''
119    return get_graph(query, max_publications, start_year, end_year, use_biobert, save_graph=save_graph, G=graph)