updating scores

app.py

@@ -6,10 +6,9 @@ import torch
 from model import MimicTransformer
 from utils import load_rule, get_attribution, get_diseases, get_drg_link, get_icd_annotations, visualize_attn
 from transformers import AutoTokenizer, AutoModel, set_seed, pipeline
-
 set_seed(42)
 model_path = 'checkpoint_0_9113.bin'
-related_tensor = torch.load('discharge_embeddings.pt')
+related_tensor = torch.nn.functional.normalize(torch.load('discharge_embeddings.pt'))
 all_summaries = pd.read_csv('all_summaries.csv')['SUMMARIES'].to_list()
 
 similarity_tokenizer = AutoTokenizer.from_pretrained('kamalkraj/BioSimCSE-BioLinkBERT-BASE')

utils.py

@@ -31,7 +31,10 @@ def get_diseases(text, pipe):
         if ent == 'B-DISEASE':
             disease_span = result['start'], result['end']
         elif ent == 'I-DISEASE':
-            disease_span = disease_span[0], result['end']
+            if len(disease_span) == 0:
+                disease_span = []
+            else:
+                disease_span = disease_span[0], result['end']
         else:
             if len(disease_span) > 1:
                 disease = text[disease_span[0]: disease_span[1]]