fix padding issue

app.py

@@ -6,7 +6,7 @@ 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')
 all_summaries = pd.read_csv('all_summaries.csv')['SUMMARIES'].to_list()
@@ -56,6 +56,7 @@ def get_model_results(text):
     logits = outputs[0][0]
     out = logits.detach().cpu()[0]
     drg_code = i2d[out.argmax().item()]
+    print(out.topk(5))
     prob = torch.nn.functional.softmax(out).max()
     return {
         'class': drg_code,

utils.py

@@ -204,7 +204,9 @@ def tokenize_icds(tokenizer, annotations, token_ids):
 
 def get_attribution(text, tokenizer, model_outputs, inputs, k=7):
     tokens = tokenizer.convert_ids_to_tokens(inputs.input_ids[0])
-    padding_idx = tokens.index('[PAD]')
+    padding_idx = 512
+    if '[PAD]' in tokens:
+        padding_idx = tokens.index('[PAD]')
     tokens = tokens[:padding_idx][1:-1]
     attn = model_outputs[-1][0]
     agg_attn, final_text = reconstruct_text(tokenizer=tokenizer, tokens=tokens, attn=attn)