dependecy parsing graph to arc standard stack operations, originally written for NLP homeword (dependency parsing)

dependecy_parsing_graph_to_arc_standard_stack_operations.py

"""
  dependecy parsing graph to arc standard stack operations
  originally written for NLP homework (dependency parsing) 
  
  WTFPL License
  by @hamidb80
  
  for more information, see "Speech and Language Processing book 3rd edition 2026" figures (19.6) and (19.7)
"""



def pick(lst, indexes):
  """
  select mulitiple indexes from a list;
  it's essentially `lst[indexes]`
  
  the name is inspired by `pick` function in BQN programming language  
  """
  return [lst[i] for i in indexes]


def list_eq(a, b):
  """
  check whether 2 lists `a` and `b` have same elements (shallow)
  """
  
  n = len(a), len(b)
  if n[0] == n[1]:
    result = all(a[i] == b[i] for i in range(n[0]))
  else:
    result = False

  if not result: print(a, "!=", b)
  return result


OPERATIONS_NAME = ("SHIFT", "LEFT-ARC", "RIGHT-ARC")

def dependency_graph_to_arc_standard(heads=[0], words=None, feature_stack_n=0, feature_buffer_n=0, log=False, ops=OPERATIONS_NAME, ROOT_TOKEN='<ROOT>'):
  """
  converts dependency graph of words in a sentence to stack operations of ARC-standard algorithm
  
  :param heads:            index of words    that corresponding word at index `i` refers to 
  :param words:            words of sentence, for debugging purposes only
  :param feature_stack_n:  number of top stack  elements to track
  :param feature_buffer_n: number of top buffer elements to track
  :param log:              show log for debugging
  :param ops:              name of operations in order of (shift, left-arc, right-arc)
  :param ROOT_TOKEN:       root token
  

  note: ROOT_TOKEN token must be present
    - `heads` must be preceded by 0
    - `words` must be preceded by 'ROOT_TOKEN'
  
  :return: tuple[operations, stack track, buffer track]
  """
  
  # ---- constants
  SHIFT_OP, LEFT_OP, RIGHT_OP = ops # name of operations may vary
  
  # ---- states
  bufferi     = [*range(len(heads))] # buffer of remaning indexes to be processed
  stacki      = []                   # stack of indexes
  operations  = []                   # list of operations that have been done
  feat_stack  = []
  feat_buffer = []
  
  dirty      = None                 # have we changed the stack? True/False

  # ---- helpers
  
  def op_track():
     feat_stack .append(stacki [-feature_stack_n:])
     feat_buffer.append(bufferi[:feature_buffer_n])
     
  def op_shift():
    if log: print(f'>> {SHIFT_OP}')

    op_track()
    operations.append(SHIFT_OP)
    stacki    .append(bufferi.pop(0))

    nonlocal dirty
    dirty = True
    
  def op_right_arc():
    if log: print(f'>> {RIGHT_OP}')

    op_track()
    operations.append(RIGHT_OP)
    stacki    .pop(-1)

    nonlocal dirty
    dirty = True

  def op_left_arc():
    if log: print(f'>> {LEFT_OP}')

    op_track()
    operations.append(LEFT_OP)
    stacki    .pop(-2)

    nonlocal dirty
    dirty = True    

  def op_log():
    print(f":: {stacki=}  {bufferi=} ::")
    
    if words:
      stackw  = pick(words,  stacki) # words[stacki]
      bufferw = pick(words, bufferi) # words[bufferi]
      print(f":: {stackw=}  {bufferw=} ::")


  # ---- algorithm
  op_shift()
  op_shift()
  
  if log: 
    assert stacki[0] == 0
    assert heads [0] == 0
    if words:
      assert words [0] == ROOT_TOKEN
    op_log()

  while stacki[-1] != 0:
    i1, i2 = stacki[-1], stacki[-2] # index aliases
    h1, h2 = heads[i1],  heads[i2]  # head  aliases
    
    future_heads = pick(heads, bufferi) # heads[bufferi]
    dirty        = False
    
    i1_free = i1 not in future_heads # is i1 free to be poped from the stack?
    i2_free = i2 not in future_heads # is i2 free to be poped from the stack?
    
    if log: # log current state
      print(f";; {future_heads=} {i1=} {i1_free=} {i2=} {i2_free=}")
    
    if not dirty and i1 == h2 and i2_free and i2 != 0: op_left_arc()
    if not dirty and i2 == h1 and i1_free:             op_right_arc()
    if not dirty:                                      op_shift()
    
    if log: # log operation
      assert dirty
      op_log()


  # ---- end
  if len(bufferi) == 0: # successful
    if log: print(f"{operations=}")
    return operations, feat_stack, feat_buffer
  else: # failed
    return None, [], []

def arc_standard_to_dependency_graph(n=0, stacks_operations=[], ops=OPERATIONS_NAME):
  """
  arc-standard stack operation to dependency graph
  
  :param n:                 number of words (including ROOT_TOKEN at the start)
  :param stacks_operations: operations
  :param ops:               name of operations in order of (shift, left-arc, right-arc)
  """
  
  SHIFT_OP, LEFT_OP, RIGHT_OP = ops
  
  bufferi = [*range(n)]
  stacki  = []
  result  = [(0,0)]
  
  for op in stacks_operations:
    if   op == SHIFT_OP:
      stacki.append(bufferi.pop(0))

    elif op == LEFT_OP:
      head = stacki    [-1]
      tail = stacki.pop(-2)
      result.append((head, tail))

    elif op == RIGHT_OP:
      head = stacki    [-2]
      tail = stacki.pop(-1)
      result.append((head, tail))
      
    else:
      raise f"invalid operation {op}"

  result.sort(key=lambda x: x[-1])
  return [r[0] for r in result]

# --------------------

eg_sentence_text1    = "<ROOT> book the flight through houston"
eg_sentence_index1   = [  0,    1,  2,     3,      4,     5]
eg_sentence_head1    = [  0,    0,  3,     1,      5,     3]
n1                   = len(eg_sentence_head1)

eg_sentence_text2    = "<ROOT> book me the morning flight"
eg_sentence_index2   = [  0,    1,  2,  3,    4,    5, ]
eg_sentence_head2    = [  0,    0,  1,  5,    5,    1, ]
n2                   = len(eg_sentence_head2)

# from "Speech and Language Processing book 3rd edition 2026" figures (19.6) and (19.7)
d1, _, _ = dependency_graph_to_arc_standard(eg_sentence_head1, eg_sentence_text1.split(), log=True, feature_stack_n=3, feature_buffer_n=2) # ['SHIFT', 'SHIFT', 'SHIFT', 'SHIFT', 'LEFT-ARC', 'SHIFT', 'SHIFT', 'LEFT-ARC', 'RIGHT-ARC', 'RIGHT-ARC', 'RIGHT-ARC']
d2, _, _ = dependency_graph_to_arc_standard(eg_sentence_head2, eg_sentence_text2.split(), log=True, feature_stack_n=3, feature_buffer_n=2) # ['SHIFT', 'SHIFT', 'SHIFT', 'RIGHT-ARC', 'SHIFT', 'SHIFT', 'SHIFT', 'LEFT-ARC', 'LEFT-ARC', 'RIGHT-ARC', 'RIGHT-ARC']

assert list_eq(arc_standard_to_dependency_graph(n1, d1), eg_sentence_head1)
assert list_eq(arc_standard_to_dependency_graph(n2, d2), eg_sentence_head2)