improved layout

app.py

@@ -52,9 +52,14 @@ except Exception as e:
 
 
 def highlight_phrase(text, phrase):
+    text = loren.fc_client.tokenizer.clean_up_tokenization(text)
     return text.replace('<mask>', f'<i><b>{phrase}</b></i>')
 
 
+def highlight_entity(text, entity):
+    return text.replace(entity, f'<i><b>{entity}</b></i>')
+
+
 def gradio_formatter(js, output_type):
     zebra_css = '''
     tr:nth-child(even) {
@@ -64,14 +69,23 @@ def gradio_formatter(js, output_type):
         background: #f1f1f1;
     }'''
     if output_type == 'e':
-        data = {'Evidence': [f'* {x}' for x in js['evidence']]}
+        data = {'Evidence': [highlight_entity(x, e) for x, e in zip(js['evidence'], js['entities'])]}
     elif output_type == 'z':
+        p_sup, p_ref, p_nei = [], [], []
+        for x in js['phrase_veracity']:
+            max_idx = torch.argmax(torch.tensor(x)).tolist()
+            x = ['%.4f' % xx for xx in x]
+            x[max_idx] = f'<i><b>{x[max_idx]}</b></i>'
+            p_sup.append(x[2])
+            p_ref.append(x[0])
+            p_nei.append(x[1])
+
         data = {
             'Claim Phrase': js['claim_phrases'],
             'Local Premise': [highlight_phrase(q, x[0]) for q, x in zip(js['cloze_qs'], js['evidential'])],
-            'p_SUP': ['%.4f' % x[2] for x in js['phrase_veracity']],
-            'p_REF': ['%.4f' % x[0] for x in js['phrase_veracity']],
-            'p_NEI': ['%.4f' % x[1] for x in js['phrase_veracity']],
+            'p_SUP': p_sup,
+            'p_REF': p_ref,
+            'p_NEI': p_nei,
         }
     else:
         raise NotImplementedError