generator.py 1.2 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849
  1. import spacy
  2. from py.error import Error
  3. nlp = spacy.load("en_core_web_trf")
  4. class Generator():
  5. """
  6. Generator class that will use a strategy resolver to generate suggestions for a text
  7. :return:
  8. """
  9. # passed text
  10. text = None
  11. # parsed text from spacy lib
  12. doc = None
  13. # the strategy resolver to use
  14. strategy_resolver = None
  15. # the parser for the presentation
  16. parser = None
  17. # suggestions calculated
  18. suggestions = None
  19. def __init__(self, text=None, strategy=None, parser=None):
  20. self.strategy_resolver = strategy
  21. self.parser = parser
  22. if text is not None:
  23. self.set_text(text)
  24. def suggest(self):
  25. if self.doc is None or self.text is None:
  26. raise Error('No document')
  27. self.suggestions = self.strategy_resolver(self.doc)
  28. return self.suggestions
  29. def set_text(self, text):
  30. self.text = text
  31. self.doc = nlp(text)
  32. def set_strategy(self, strategy_resolver):
  33. self.strategy_resolver = strategy_resolver
  34. def present(self):
  35. self.parser.suggestions = self.suggestions
  36. return self.parser.present()