leaf_strategy.py 639 B

1234567891011121314151617181920212223242526272829
  1. from spacy.tokens import Token
  2. from word_processor.types import DEP_TYPES
  3. def leaf_strategy(doc) -> [[Token]]:
  4. """
  5. Should return an arrays of variable names based on leaf strategy
  6. :param doc: spacy document
  7. :return Array of strings
  8. """
  9. for token in doc:
  10. if token.dep_ == DEP_TYPES['ROOT']:
  11. suggestions = dfs(token)
  12. break
  13. return suggestions or []
  14. def dfs(graph, result=[], output=[]):
  15. flag = False
  16. for u in graph.children:
  17. flag = True
  18. dfs(u, [*result, graph], output)
  19. if flag is False:
  20. output.append([*result, graph])
  21. return output