Update app.py

app.py

@@ -9,40 +9,32 @@ from models.spabert.utils.common_utils import load_spatial_bert_pretrained_weigh
 
 from PIL import Image
 
-##LOAD MODEL SECTION##
+device = torch.device('cpu')
+
+#Spacy Initialization Section
 nlp = spacy.load("./models/en_core_web_sm")
 
 #BERT Initialization Section
 bert_tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
 bert_model = BertModel.from_pretrained("bert-base-uncased")
+bert_model.to(device)
 bert_model.eval()
 
 #SpaBERT Initialization Section
 data_file_path = 'models/spabert/datasets/SPABERT_finetuning_data_combined.json'
-pretrained_model = 'models/spabert/datasets/fine-spabert-base-uncased-finetuned-osm-mn.pth'
+pretrained_model_path = 'models/spabert/datasets/fine-spabert-base-uncased-finetuned-osm-mn.pth'
+pre_trained_model = torch.load(pretrained_model_path, map_location=torch.device('cpu'))
 
-b_model =  BertForMaskedLM.from_pretrained('bert-base-uncased')
-b_tokenizer =  BertTokenizer.from_pretrained('bert-base-uncased')
+#b_model =  BertForMaskedLM.from_pretrained('bert-base-uncased')
+#b_tokenizer =  BertTokenizer.from_pretrained('bert-base-uncased')
 
 config = SpatialBertConfig()
 config.output_hidden_states = True
-
 spaBERT_model = SpatialBertForMaskedLM(config)
-spaBERT_model.load_state_dict(b_model.state_dict(), strict = False)
 
-pre_trained_model = torch.load(pretrained_model, map_location=torch.device('cpu'))
-
-model_keys = spaBERT_model.state_dict()
-cnt_layers = 0
-for key in model_keys:
-    if key in pre_trained_model:
-        model_keys[key] = pre_trained_model[key]
-        cnt_layers += 1
-    else:
-        print("No weight for", key)
-print(cnt_layers, 'layers loaded')
+spaBERT_model.load_state_dict(bert_model.state_dict(), strict = False)
+spaBERT_model.load_state_dict(pre_trained_model, strict=False)
 
-spaBERT_model.load_state_dict(model_keys)
 spaBERT_model.to(device)
 spaBERT_model.eval()
 
@@ -69,7 +61,7 @@ COLOR_MAP = {
     'GPE': ('green', 'Geopolitical Entities (e.g., countries, cities)')
 }
 
-# Display the color key with descriptions
+# Display the color key
 st.write("**Color Key:**")
 for label, (color, description) in COLOR_MAP.items():
     st.markdown(f"- **{label}**: <span style='color:{color}'>{color}</span> - {description}", unsafe_allow_html=True)