Update app.py

app.py

@@ -6,6 +6,8 @@ from transformers.models.bert.modeling_bert import BertForMaskedLM
 
 from models.spabert.models.spatial_bert_model import SpatialBertConfig, SpatialBertForMaskedLM, SpatialBertModel
 from models.spabert.utils.common_utils import load_spatial_bert_pretrained_weights
+from models.spabert.datasets.osm_sample_loader import PbfMapDataset
+from torch.utils.data import DataLoader
 
 from PIL import Image
 
@@ -35,6 +37,19 @@ spaBERT_model.load_state_dict(pre_trained_model, strict=False)
 spaBERT_model.to(device)
 spaBERT_model.eval()
 
+# Load data using SpatialDataset
+spatialDataset = PbfMapDataset(data_file_path = data_file_path,
+                                        tokenizer = tokenizer,
+                                        #max_token_len = 256,              #Originally 300
+                                        max_token_len = max_seq_length,              #Originally 300
+                                        distance_norm_factor = 0.0001,
+                                        spatial_dist_fill = 20,
+                                        with_type = False,
+                                        sep_between_neighbors = True,    #Initially false, play around with this potentially?
+                                        label_encoder = None,             #Initially None, potentially change this because we do have real/fake reviews.
+                                        mode = None)                      #If set to None it will use the full dataset for mlm
+
+data_loader = DataLoader(spatialDataset, batch_size=1, num_workers=0, shuffle=False, pin_memory=False, drop_last=False) #issue needs to be fixed with num_workers not stopping after finished
 
 #Get BERT Embedding for review
 def get_bert_embedding(review_text):