code: def _make_probe_inputs(
      edges: list[EdgeLabel],
      embs: np.ndarray,
      offsets: np.ndarray,
      src_idx: int,
  ):
      for j, edge in enumerate(edges):
          span1_embs = _extract_span(edge.span1, embs, offsets)
          span2_embs = _extract_span(edge.span2, embs, offsets)
          yield {
              "span1_embs": span1_embs,
              "span2_embs": span2_embs,
              "label": edge.label,
              "src_idx": src_idx,
              "edge_idx": j,
          }
  
imports: import os
  from collections.abc import Iterable
  from typing import Optional
  from absl import logging
  from lit_nlp.api import dataset
  from lit_nlp.api import dtypes
  from lit_nlp.api import model
  from lit_nlp.api import types
  from lit_nlp.lib import utils
  import numpy
  import tensorflow
label_desc: Use image node_id 2 for calling a global function with example usage: _make_probe_inputs(edges, embs, offsets, src_idx) without return types