瀏覽代碼

make generator

Nikos Atlas 2 年之前
父節點
當前提交
060bfad8b5
共有 1 個文件被更改,包括 35 次插入0 次删除
  1. 35 0
      src/generators/generator.py

+ 35 - 0
src/generators/generator.py

@@ -0,0 +1,35 @@
+import spacy
+from py.error import Error
+
+nlp = spacy.load("en_core_web_trf")
+
+
+class Generator():
+    """
+        Generator class that will use a strategy resolver to generate suggestions for a text
+        :return:
+
+    """
+    text = None
+    doc = None
+    strategy_resolver = None
+
+    def __init__(self, text=None, strategy=None):
+        self.strategy_resolver = strategy
+        if text is not None:
+            self.set_text(text)
+
+    def suggest(self):
+        if self.doc is None or self.text is None:
+            raise Error('No document')
+
+        suggestions = self.strategy_resolver(self.doc)
+
+        return suggestions
+
+    def set_text(self, text):
+        self.text = text
+        self.doc = nlp(text)
+
+    def set_strategy(self, strategy_resolver):
+        self.strategy_resolver = strategy_resolver