Update tree.py

tree.py

@@ -5,18 +5,17 @@ from collections import defaultdict
 from paraphraser import generate_paraphrase
 from masking_methods import mask, mask_non_stopword
 
-def generate_plot(original_sentence):
-    paraphrased_sentences = generate_paraphrase(original_sentence)
-    first_paraphrased_sentence = paraphrased_sentences[0]
+def generate_plot(original_sentence, selected_sentences):
+    first_paraphrased_sentence = selected_sentences[0]
     masked_sentence = mask_non_stopword(first_paraphrased_sentence)
     masked_versions = mask(masked_sentence)
     
     nodes = []
     nodes.append(original_sentence)
-    nodes.extend(paraphrased_sentences)
+    nodes.extend(selected_sentences)
     nodes.extend(masked_versions)
     nodes[0] += ' L0'
-    para_len = len(paraphrased_sentences)
+    para_len = len(selected_sentences)
     for i in range(1, para_len+1):
         nodes[i] += ' L1'
     for i in range(para_len+1, len(nodes)):