code: def extract_2hop_graph(
      in_context_ents, kg_graph
  ) -> Dict[Any, Dict[Any, bool]]:
      """For each wikipedia page with N in-context entities, extract a subgraph that contains only 2hop paths between any pair of nodes.
  
      Args:
        in_context_ents: all entities within each wiki-page, stored as dict.
        kg_graph: the global KG (i.e. WikiData Knowledge Graph), stored as dict.
  
      Returns:
        Extracted 2-hop subgraph for each page.
      """
      all_nodes = {se: [se] for se in in_context_ents}
      for se in in_context_ents:
          if se in kg_graph:
              for te in kg_graph[se]:
                  if te not in in_context_ents:
                      if te not in all_nodes:
                          all_nodes[te] = [se]
                      else:
                          all_nodes[te] += [se]
      remain_nodes = {
          e: True for e in all_nodes if len(all_nodes[e]) > 1
      }
      for e in in_context_ents:
          remain_nodes[e] = True
  
      two_graph = {}
      for se in remain_nodes:
          if se in kg_graph:
              for te in kg_graph[se]:
                  if te in remain_nodes:
                      if se not in two_graph:
                          two_graph[se] = {}
                      two_graph[se][te] = kg_graph[se][te]
                      if te in kg_graph and se in kg_graph[te]:
                          if te not in two_graph:
                              two_graph[te] = {}
                          two_graph[te][se] = kg_graph[te][se]
      return two_graph
  
file_path: scenic/scenic/projects/knowledge_visual_language/data/wikidata/data_util.py
imports: import json
  from typing import Any, Dict, List
  import matplotlib.pyplot
  import networkx
  import numpy
  import tqdm
  import tensorflow
  from tf.io import gfile
label_desc: Use image node_id 3 for calling a global function with example usage: extract_2hop_graph(in_context_ents, kg_graph) and returns: two_graph