docstr: Produce initial word queues out of the model's LSTMs for use in the tagged word lists.
  
  Operates in a batched fashion to reduce the runtime for the LSTM operations
code: def initial_word_queues(self, tagged_word_lists):
      """
      Produce initial word queues out of the model's LSTMs for use in the tagged word lists.
  
      Operates in a batched fashion to reduce the runtime for the LSTM operations
      """
      device = next(self.parameters()).device
  
      vocab_map = self.vocab_map
  
      def map_word(word):
          idx = vocab_map.get(word, None)
          if idx is not None:
              return idx
          return vocab_map.get(word.lower(), UNK_ID)
  
      all_word_inputs = []
      all_word_labels = [
          [word.children[0].label for word in tagged_words]
          for tagged_words in tagged_word_lists
      ]
  
      for sentence_idx, tagged_words in enumerate(tagged_word_lists):
          word_labels = all_word_labels[sentence_idx]
          word_idx = torch.stack(
              [
                  self.vocab_tensors[map_word(word.children[0].label)]
                  for word in tagged_words
              ]
          )
          word_input = self.embedding(word_idx)
  
          # this occasionally learns UNK at train time
          if self.training:
              delta_labels = [
                  None
                  if word in self.rare_words
                  and random.random()
                  < self.args["rare_word_unknown_frequency"]
                  else word
                  for word in word_labels
              ]
          else:
              delta_labels = word_labels
          delta_idx = torch.stack(
              [
                  self.delta_tensors[
                      self.delta_word_map.get(word, UNK_ID)
                  ]
                  for word in delta_labels
              ]
          )
  
          delta_input = self.delta_embedding(delta_idx)
          word_inputs = [word_input, delta_input]
  
          if self.tag_embedding_dim > 0:
              if self.training:
                  tag_labels = [
                      None
                      if random.random()
                      < self.args["tag_unknown_frequency"]
                      else word.label
                      for word in tagged_words
                  ]
              else:
                  tag_labels = [word.label for word in tagged_words]
              tag_idx = torch.stack(
                  [
                      self.tag_tensors[self.tag_map.get(tag, UNK_ID)]
                      for tag in tag_labels
                  ]
              )
              tag_input = self.tag_embedding(tag_idx)
              word_inputs.append(tag_input)
  
          all_word_inputs.append(word_inputs)
  
      if self.forward_charlm is not None:
          all_forward_chars = (
              self.forward_charlm.build_char_representation(
                  all_word_labels
              )
          )
          for word_inputs, forward_chars in zip(
              all_word_inputs, all_forward_chars
          ):
              word_inputs.append(forward_chars)
      if self.backward_charlm is not None:
          all_backward_chars = (
              self.backward_charlm.build_char_representation(
                  all_word_labels
              )
          )
          for word_inputs, backward_chars in zip(
              all_word_inputs, all_backward_chars
          ):
              word_inputs.append(backward_chars)
  
      all_word_inputs = [
          torch.cat(word_inputs, dim=1)
          for word_inputs in all_word_inputs
      ]
      if self.sentence_boundary_vectors is not SentenceBoundary.NONE:
          word_start = self.word_start_embedding.unsqueeze(0)
          word_end = self.word_end_embedding.unsqueeze(0)
          all_word_inputs = [
              torch.cat([word_start, word_inputs, word_end], dim=0)
              for word_inputs in all_word_inputs
          ]
  
      if self.bert_model is not None:
          # BERT embedding extraction
          # result will be len+2 for each sentence
          # we will take 1:-1 if we don't care about the endpoints
          bert_embeddings = extract_bert_embeddings(
              self.args["bert_model"],
              self.bert_tokenizer,
              self.bert_model,
              all_word_labels,
              device,
              keep_endpoints=self.sentence_boundary_vectors
              is not SentenceBoundary.NONE,
              num_layers=self.bert_layer_mix.in_features
              if self.bert_layer_mix is not None
              else None,
              detach=not self.args["bert_finetune"]
              and not self.args["stage1_bert_finetune"],
          )
          if self.bert_layer_mix is not None:
              # add the average so that the default behavior is to
              # take an average of the N layers, and anything else
              # other than that needs to be learned
              bert_embeddings = [
                  self.bert_layer_mix(feature).squeeze(2)
                  + feature.sum(axis=2)
                  / self.bert_layer_mix.in_features
                  for feature in bert_embeddings
              ]
  
          all_word_inputs = [
              torch.cat((x, y), axis=1)
              for x, y in zip(all_word_inputs, bert_embeddings)
          ]
  
      # Extract partitioned representation
      if self.partitioned_transformer_module is not None:
          partitioned_embeddings = self.partitioned_transformer_module(
              None, all_word_inputs
          )
          all_word_inputs = [
              torch.cat((x, y[: x.shape[0], :]), axis=1)
              for x, y in zip(all_word_inputs, partitioned_embeddings)
          ]
  
      # Extract Labeled Representation
      if self.label_attention_module is not None:
          if self.args["lattn_combined_input"]:
              labeled_representations = self.label_attention_module(
                  all_word_inputs, tagged_word_lists
              )
          else:
              labeled_representations = self.label_attention_module(
                  partitioned_embeddings, tagged_word_lists
              )
          all_word_inputs = [
              torch.cat((x, y[: x.shape[0], :]), axis=1)
              for x, y in zip(all_word_inputs, labeled_representations)
          ]
  
      all_word_inputs = [
          self.word_dropout(word_inputs)
          for word_inputs in all_word_inputs
      ]
      packed_word_input = torch.nn.utils.rnn.pack_sequence(
          all_word_inputs, enforce_sorted=False
      )
      word_output, _ = self.word_lstm(packed_word_input)
      # would like to do word_to_constituent here, but it seems PackedSequence doesn't support Linear
      # word_output will now be sentence x batch x 2*hidden_size
      (
          word_output,
          word_output_lens,
      ) = torch.nn.utils.rnn.pad_packed_sequence(word_output)
      # now sentence x batch x hidden_size
  
      word_queues = []
      for sentence_idx, tagged_words in enumerate(tagged_word_lists):
          if (
              self.sentence_boundary_vectors
              is not SentenceBoundary.NONE
          ):
              sentence_output = word_output[
                  : len(tagged_words) + 2, sentence_idx, :
              ]
          else:
              sentence_output = word_output[
                  : len(tagged_words), sentence_idx, :
              ]
          sentence_output = self.word_to_constituent(sentence_output)
          sentence_output = self.nonlinearity(sentence_output)
          # TODO: this makes it so constituents downstream are
          # build with the outputs of the LSTM, not the word
          # embeddings themselves.  It is possible we want to
          # transform the word_input to hidden_size in some way
          # and use that instead
          if (
              self.sentence_boundary_vectors
              is not SentenceBoundary.NONE
          ):
              word_queue = [WordNode(None, sentence_output[0, :])]
              word_queue += [
                  WordNode(tag_node, sentence_output[idx + 1, :])
                  for idx, tag_node in enumerate(tagged_words)
              ]
              word_queue.append(
                  WordNode(
                      None, sentence_output[len(tagged_words) + 1, :]
                  )
              )
          else:
              word_queue = [WordNode(None, self.word_zeros)]
              word_queue += [
                  WordNode(tag_node, sentence_output[idx, :])
                  for idx, tag_node in enumerate(tagged_words)
              ]
              word_queue.append(WordNode(None, self.word_zeros))
  
          if self.reverse_sentence():
              word_queue = list(reversed(word_queue))
          word_queues.append(word_queue)
  
      return word_queues
  
imports: from collections import namedtuple
  from enum import Enum
  import logging
  import math
  import random
  import torch
  import torch.nn
  from torch.nn.utils.rnn import pack_padded_sequence
  from stanza.models.common.bert_embedding import extract_bert_embeddings
  from stanza.models.common.maxout_linear import MaxoutLinear
  from stanza.models.common.utils import unsort
  from stanza.models.common.vocab import PAD_ID, UNK_ID
  from stanza.models.constituency.base_model import BaseModel
  from stanza.models.constituency.label_attention import LabelAttentionModule
  from stanza.models.constituency.lstm_tree_stack import LSTMTreeStack
  from stanza.models.constituency.parse_transitions import TransitionScheme
  from stanza.models.constituency.parse_tree import Tree
  from stanza.models.constituency.partitioned_transformer import PartitionedTransformerModule
  from stanza.models.constituency.positional_encoding import ConcatSinusoidalEncoding
  from stanza.models.constituency.transformer_tree_stack import TransformerTreeStack
  from stanza.models.constituency.tree_stack import TreeStack
  from stanza.models.constituency.utils import build_nonlinearity, initialize_linear
label_desc: Use image node_id 15 for calling the LSTMModel obj's underlying member method code with example usage: obj.initial_word_queues(tagged_word_lists) and returns: word_queues, vocab_map, idx