fixed model issue

all_summaries.csv

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:3c74c03c1a4b5ad01eff9eea8a7062c660d342bf912794191cd5e3ebeb4abe44
-size 1114819287
+oid sha256:64b8415ab96d12e61393c8ca12ebb844ab32a57df314984e98e91e1064bebf41
+size 640698121

app.py

@@ -4,9 +4,14 @@ import pandas as pd
 import torch
 
 from model import MimicTransformer
-from utils import load_rule, get_attribution, get_diseases, get_drg_link, get_icd_annotations, visualize_attn
+from utils import load_rule, get_attribution, get_diseases, get_drg_link, get_icd_annotations, visualize_attn, clean_text
 from transformers import AutoTokenizer, AutoModel, set_seed, pipeline
 
+torch.manual_seed(0)
+set_seed(34)
+if torch.cuda.is_available():
+    torch.cuda.manual_seed_all(0)
+
 model_path = 'checkpoint_0_9113.bin'
 related_tensor = torch.load('discharge_embeddings.pt')
 all_summaries = pd.read_csv('all_summaries.csv')['SUMMARIES'].to_list()
@@ -16,7 +21,9 @@ similarity_model = AutoModel.from_pretrained('kamalkraj/BioSimCSE-BioLinkBERT-BA
 similarity_model.eval()
 
 def read_model(model, path):
-    model.load_state_dict(torch.load(path, map_location=torch.device('cpu')), strict=False)
+    # model.load_state_dict(torch.load(path, map_location=torch.device('cpu')))
+    state_dict = torch.load(path, map_location='cpu')
+    model.load_state_dict({"model."+k: v for k, v in state_dict.items()}, strict=False)
     return model
 
 mimic = MimicTransformer(cutoff=512)
@@ -50,8 +57,10 @@ def mean_pooling(model_output, attention_mask):
 
 
 def get_model_results(text):
+    text = clean_text(text)
     inputs = tokenizer(text, return_tensors='pt', padding='max_length', max_length=512, truncation=True)
-    outputs = mimic(input_ids=inputs.input_ids, attention_mask=inputs.attention_mask, drg_labels=None)
+    with torch.no_grad():
+        outputs = mimic(input_ids=inputs.input_ids, attention_mask=inputs.attention_mask, drg_labels=None)
     attribution, reconstructed_text = get_attribution(text=text, tokenizer=tokenizer, model_outputs=outputs, inputs=inputs, k=10)
     logits = outputs[0][0]
     out = logits.detach().cpu()[0]
@@ -93,7 +102,12 @@ def find_related_summaries(text):
 
 
 def run(text, related_discharges=False):
+    torch.manual_seed(0)
+    set_seed(34)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(0)
     # initial drg results
+
     model_results = get_model_results(text=text)
     drg_code = model_results['class']
 

discharge_embeddings.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:2179abc7e448f3bb4a091f4f27eadc5b1d4829464eeb9bfb7fd9c6363844aaaa
-size 1228800786
+oid sha256:cbc05e83aa36756a35bee2f104e3c3dcc8fb1f26442d89ff52916d7052cd036b
+size 713869074

model.py

@@ -1,4 +1,4 @@
-from transformers import AutoConfig, AutoTokenizer, AutoModelForSequenceClassification
+from transformers import AutoConfig, AutoTokenizer, AutoModelForSequenceClassification, set_seed
 from torch.utils.data import DataLoader
 from torch.nn import Linear, Module
 from typing import Dict, List
@@ -6,6 +6,11 @@ from collections import Counter, defaultdict
 from itertools import chain
 import torch
 
+torch.manual_seed(0)
+set_seed(34)
+if torch.cuda.is_available():
+    torch.cuda.manual_seed_all(0)
+
 class MimicTransformer(Module):
     def __init__(self, num_labels=738, tokenizer_name='clinical', cutoff=512):
         """
@@ -17,6 +22,7 @@ class MimicTransformer(Module):
         self.config = AutoConfig.from_pretrained(self.tokenizer_name, num_labels=self.num_labels)
         self.tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_name, config=self.config)
         self.model = AutoModelForSequenceClassification.from_pretrained(self.tokenizer_name, config=self.config)
+        self.model.eval()
         if 'longformer' in self.tokenizer_name:
             self.cutoff = self.model.config.max_position_embeddings
         else: