Update inductor.py

inductor.py

@@ -82,17 +82,23 @@ class BartInductor(object):
         else:
             return text
 
-    def generate(self, inputs, k=10, topk=10):
+    def generate(self, inputs, k=10, topk=10, return_scores=False):
         with torch.no_grad():
             tB_probs = self.generate_rule(inputs, k)
-            #ret = [t[0].replace('<ent0>','<mask>').replace('<ent1>','<mask>') for t in tB_probs]
-            ret = [(t[0].replace('<ent0>','<mask>').replace('<ent1>','<mask>'), t[1]) for t in tB_probs]
-
-            new_ret = []
-            for temp in ret:
-                temp = (self.clean(temp[0].strip()), temp[1])
-                if len(new_ret) < topk and temp not in new_ret:
-                    new_ret.append(temp)
+            if return_scores:
+                ret = [(t[0].replace('<ent0>','<mask>').replace('<ent1>','<mask>'), t[1]) for t in tB_probs]
+                new_ret = []
+                for temp in ret:
+                    temp = (self.clean(temp[0].strip()), temp[1])
+                    if len(new_ret) < topk and temp not in new_ret:
+                        new_ret.append(temp)
+            else:
+                ret = [t[0].replace('<ent0>','<mask>').replace('<ent1>','<mask>') for t in tB_probs]
+                new_ret = []
+                for temp in ret:
+                    temp = self.clean(temp[0].strip())
+                    if len(new_ret) < topk and temp not in new_ret:
+                        new_ret.append(temp)
 
             return new_ret