Upload inductor.py

inductor.py

@@ -76,25 +76,25 @@ class BartInductor(object):
         self.stop_weight = stop_weight[0, :]
 
     def clean(self, text):
-        segments = text.split('<mask>')
-        if len(segments) == 3 and segments[2].startswith('.'):
-            return '<mask>'.join(segments[:2]) + '<mask>.'
+        segments = re.split(r'<ent\d>', text)
+        last_segment = segments[-1]
+        if last_segment.startswith('.'):
+            return text[:text.rfind(last_segment)]+'.'
         else:
             return text
 
     def generate(self, inputs, k=10, topk=10, return_scores=False):
         with torch.no_grad():
             tB_probs = self.generate_rule(inputs, k)
+            new_ret = []
             if return_scores:
-                ret = [(t[0].replace('<ent0>','<mask>').replace('<ent1>','<mask>'), t[1]) for t in tB_probs]
-                new_ret = []
+                ret = [(t[0], t[1]) for t in tB_probs]
                 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 = []
+                ret = [t[0] for t in tB_probs]
                 for temp in ret:
                     temp = self.clean(temp.strip())
                     if len(new_ret) < topk and temp not in new_ret:
@@ -134,7 +134,7 @@ class BartInductor(object):
         return ret
 
     def extract_words_for_tA_bart(self, tA, k=6, print_it = False):
-        spans = [t.lower().strip() for t in tA[:-1].split('<mask>')]
+        spans = [t.lower().strip() for t in re.split(r'<.*?>', tA[:-1])]
         generated_ids = self.tokenizer([tA], padding='longest', return_tensors='pt')['input_ids'].to(device).to(torch.int64)
         generated_ret = self.orion_instance_generator.generate(generated_ids, num_beams=max(120, k),
                                             #num_beam_groups=max(120, k),
@@ -300,6 +300,7 @@ class BartInductor(object):
         for k1 in tB_prob:
             ret.append([k1, tB_prob[k1]])
         ret = sorted(ret, key=lambda x: x[1], reverse=True)[:k]
+
         if self.if_then:
             for i, temp in enumerate(ret):
                 sentence = temp[0]