Upload 76 files

.gitattributes

@@ -37,3 +37,7 @@ models/en_core_web_sm/en_core_web_sm-3.7.1/ner/model filter=lfs diff=lfs merge=l
 models/en_core_web_sm/en_core_web_sm-3.7.1/tok2vec/model filter=lfs diff=lfs merge=lfs -text
 models/en_core_web_sm/ner/model filter=lfs diff=lfs merge=lfs -text
 models/en_core_web_sm/tok2vec/model filter=lfs diff=lfs merge=lfs -text
+models/spabert/notebooks/tutorial_datasets/output.csv.json filter=lfs diff=lfs merge=lfs -text
+models/spabert/notebooks/tutorial_datasets/spabert_osm_mn.json filter=lfs diff=lfs merge=lfs -text
+models/spabert/notebooks/tutorial_datasets/spabert_whg_wikidata.json filter=lfs diff=lfs merge=lfs -text
+models/spabert/notebooks/tutorial_datasets/spabert_wikidata_sampled.json filter=lfs diff=lfs merge=lfs -text

models/spabert/README.md

@@ -0,0 +1,69 @@
+# SpaBERT: A Pretrained Language Model from Geographic Data for Geo-Entity Representation
+
+This repo contains code for [SpaBERT: A Pretrained Language Model from Geographic Data for Geo-Entity Representation](https://arxiv.org/abs/2210.12213) which was published in EMNLP 2022. SpaBERT provides a general-purpose geo-entity representation based on neighboring entities in geospatial data. SpaBERT extends BERT to capture linearized spatial context, while incorporating a spatial coordinate embedding mechanism to preserve spatial relations of entities in the 2-dimensional space. SpaBERT is pretrained with masked language modeling and masked entity prediction tasks to learn spatial dependencies.
+
+* Slides: [emnlp22-spabert.pdf](https://drive.google.com/file/d/1V1URsRfpw13dbkb_zgBXeNqZJ0AF2744/view?usp=share_link)
+
+
+## Pretraining 
+Pretrained model weights can be downloaded from the Google Drive for [SpaBERT-base](https://drive.google.com/file/d/1l44FY3DtDxzM_YVh3RR6PJwKnl80IYWB/view?usp=sharing) and [SpaBERT-large](https://drive.google.com/file/d/1LeZayTR92R5bu9gH_cGCwef7nnMX35cR/view?usp=share_link). 
+
+Weights can also obtained from training from scratch using the following sample code. Data for pretraining can be downloaded [here](https://drive.google.com/drive/folders/1eaeVvUCcJVcNwnyTCk-1N1IKfukihk4j?usp=share_link). 
+
+* Code to pretrain SpaBERT-base model:
+
+  ```python3 train_mlm.py --lr=5e-5 --sep_between_neighbors --bert_option='bert-base'```
+
+* Code to pretrain SpaBERT-large model:
+
+  ```python3 train_mlm.py --lr=1e-6 --sep_between_neighbors --bert_option='bert-large```
+  
+## Downstream Tasks
+### Supervised Geo-entity typing 
+The goal is to predict a geo-entity’s semantic type (e.g., transportation and healthcare) given the target geo-entity name and spatial context (i.e. surrounding neighbors name and location). 
+
+Models trained on OSM in London and California region can be downloaded from Google Drive for [SpaBERT-base](https://drive.google.com/file/d/1XFcA3sxC4wTlt7VjvMp1zNrWY5rjafzE/view?usp=share_link) and [SpaBERT-large](https://drive.google.com/file/d/12_FDVeSYkl_HQ61JmuMU6cRjQdKNpgR_/view?usp=share_link)
+
+Data used for training and testing can be downloaded [here](https://drive.google.com/drive/folders/1uyvGdiJdu-Cym4dOKhQLIkKpfgHvfo01?usp=share_link)
+
+* Sample code for training SpaBERT-base typing model
+
+```
+python3 train_cls_spatialbert.py --lr=5e-5 --sep_between_neighbors --bert_option='bert-base'  --with_type --mlm_checkpoint_path='mlm_mem_keeppos_ep0_iter06000_0.2936.pth' 
+```
+
+* Sample code for training SpaBERT-large typing model
+
+```
+python3 train_cls_spatialbert.py --lr=1e-6 --sep_between_neighbors --bert_option='bert-large'  --with_type --mlm_checkpoint_path='mlm_mem_keeppos_ep1_iter02000_0.4400.pth' --epochs=20
+```
+
+### Unsupervised Geo-entity Linking
+
+Geo-entity linking is to link geo-entities from a geographic information system (GIS) oriented dataset to a knowledge base (KB). This task unsupervised thus does not require any further training. Pretrained models can be directly used for this task. 
+
+
+Linking with SpaBERT-base
+```
+python3 unsupervised_wiki_location_allcand.py --model_name='spatial_bert-base' --sep_between_neighbors \
+ --spatial_bert_weight_dir='weights/' --spatial_bert_weight_name='mlm_mem_keeppos_ep0_iter06000_0.2936.pth'
+
+```
+
+Linking with SpaBERT-large
+```
+python3 unsupervised_wiki_location_allcand.py --model_name='spatial_bert-large' --sep_between_neighbors \
+ --spatial_bert_weight_dir='weights/' --spatial_bert_weight_name='mlm_mem_keeppos_ep1_iter02000_0.4400.pth'
+```
+
+Data used for linking from USGS historical maps to WikiData KB is provided [here](https://drive.google.com/drive/folders/1qKJnj71qxnca_TaygK-Y3EIySnMyFpFn?usp=share_link)
+
+## Acknowledgement
+```
+@article{li2022spabert,
+  title={SpaBERT: A Pretrained Language Model from Geographic Data for Geo-Entity Representation},
+  author={Zekun Li, Jina Kim, Yao-Yi Chiang and Muhao Chen},
+  journal={EMNLP},
+  year={2022}
+}
+```

models/spabert/datasets/const.py

@@ -0,0 +1,162 @@
+def revert_dict(coarse_to_fine_dict):
+    fine_to_coarse_dict = dict()
+    for key, value in coarse_to_fine_dict.items():
+        for v in value:
+            fine_to_coarse_dict[v] = key
+    return fine_to_coarse_dict
+
+CLASS_9_LIST = ['education', 'entertainment_arts_culture', 'facilities', 'financial', 'healthcare', 'public_service', 'sustenance', 'transportation', 'waste_management']
+
+CLASS_118_LIST=['animal_boarding', 'animal_breeding', 'animal_shelter', 'arts_centre', 'atm', 'baby_hatch', 'baking_oven', 'bank', 'bar', 'bbq', 'bench', 'bicycle_parking', 'bicycle_rental', 'bicycle_repair_station', 'biergarten', 'boat_rental', 'boat_sharing', 'brothel', 'bureau_de_change', 'bus_station', 'cafe', 'car_rental', 'car_sharing', 'car_wash', 'casino', 'charging_station', 'childcare', 'cinema', 'clinic', 'clock', 'college', 'community_centre', 'compressed_air', 'conference_centre', 'courthouse', 'crematorium', 'dentist', 'dive_centre', 'doctors', 'dog_toilet', 'dressing_room', 'drinking_water', 'driving_school', 'events_venue', 'fast_food', 'ferry_terminal', 'fire_station', 'food_court', 'fountain', 'fuel', 'funeral_hall', 'gambling', 'give_box', 'grave_yard', 'grit_bin', 'hospital', 'hunting_stand', 'ice_cream', 'internet_cafe', 'kindergarten', 'kitchen', 'kneipp_water_cure', 'language_school', 'library', 'lounger', 'love_hotel', 'marketplace', 'monastery', 'motorcycle_parking', 'music_school', 'nightclub', 'nursing_home', 'parcel_locker', 'parking', 'parking_entrance', 'parking_space', 'pharmacy', 'photo_booth', 'place_of_mourning', 'place_of_worship', 'planetarium', 'police', 'post_box', 'post_depot', 'post_office', 'prison', 'pub', 'public_bath', 'public_bookcase', 'ranger_station', 'recycling', 'refugee_site', 'restaurant', 'sanitary_dump_station', 'school', 'shelter', 'shower', 'social_centre', 'social_facility', 'stripclub', 'studio', 'swingerclub', 'taxi', 'telephone', 'theatre', 'toilets', 'townhall', 'toy_library', 'training', 'university', 'vehicle_inspection', 'vending_machine', 'veterinary', 'waste_basket', 'waste_disposal', 'waste_transfer_station', 'water_point', 'watering_place']
+
+CLASS_95_LIST = ['arts_centre', 'atm', 'baby_hatch', 'bank', 'bar', 'bbq', 'bench', 'bicycle_parking', 'bicycle_rental', 'bicycle_repair_station', 'biergarten', 'boat_rental', 'boat_sharing', 'brothel', 'bureau_de_change', 'bus_station', 'cafe', 'car_rental', 'car_sharing', 'car_wash', 'casino', 'charging_station', 'cinema', 'clinic', 'college', 'community_centre', 'compressed_air', 'conference_centre', 'courthouse', 'dentist', 'doctors', 'dog_toilet', 'dressing_room', 'drinking_water', 'driving_school', 'events_venue', 'fast_food', 'ferry_terminal', 'fire_station', 'food_court', 'fountain', 'fuel', 'gambling', 'give_box', 'grit_bin', 'hospital', 'ice_cream', 'kindergarten', 'language_school', 'library', 'love_hotel', 'motorcycle_parking', 'music_school', 'nightclub', 'nursing_home', 'parcel_locker', 'parking', 'parking_entrance', 'parking_space', 'pharmacy', 'planetarium', 'police', 'post_box', 'post_depot', 'post_office', 'prison', 'pub', 'public_bookcase', 'ranger_station', 'recycling', 'restaurant', 'sanitary_dump_station', 'school', 'shelter', 'shower', 'social_centre', 'social_facility', 'stripclub', 'studio', 'swingerclub', 'taxi', 'telephone', 'theatre', 'toilets', 'townhall', 'toy_library', 'training', 'university', 'vehicle_inspection', 'veterinary', 'waste_basket', 'waste_disposal', 'waste_transfer_station', 'water_point', 'watering_place']
+
+CLASS_74_LIST = ['arts_centre', 'atm', 'bank', 'bar', 'bench', 'bicycle_parking', 'bicycle_rental', 'bicycle_repair_station', 'boat_rental', 'bureau_de_change', 'bus_station', 'cafe', 'car_rental', 'car_sharing', 'car_wash', 'charging_station', 'cinema', 'clinic', 'college', 'community_centre', 'conference_centre', 'courthouse', 'dentist', 'doctors', 'drinking_water', 'driving_school', 'events_venue', 'fast_food', 'ferry_terminal', 'fire_station', 'food_court', 'fountain', 'fuel', 'gambling', 'hospital', 'kindergarten', 'language_school', 'library', 'motorcycle_parking', 'music_school', 'nightclub', 'nursing_home', 'parcel_locker', 'parking', 'pharmacy', 'police', 'post_box', 'post_depot', 'post_office', 'pub', 'public_bookcase', 'recycling', 'restaurant', 'sanitary_dump_station', 'school', 'shelter', 'social_centre', 'social_facility', 'stripclub', 'studio', 'swingerclub', 'taxi', 'telephone', 'theatre', 'toilets', 'townhall', 'university', 'vehicle_inspection', 'veterinary', 'waste_basket', 'waste_disposal', 'waste_transfer_station', 'water_point', 'watering_place']
+
+FEWSHOT_CLASS_55_LIST = ['arts_centre', 'atm', 'bank', 'bar', 'bench', 'bicycle_parking', 'bicycle_rental', 'boat_rental', 'bureau_de_change', 'bus_station', 'cafe', 'car_rental', 'car_sharing', 'car_wash', 'charging_station', 'cinema', 'clinic', 'college', 'community_centre', 'courthouse', 'dentist', 'doctors', 'drinking_water', 'driving_school', 'events_venue', 'fast_food', 'ferry_terminal', 'fire_station', 'fountain', 'fuel', 'hospital', 'kindergarten', 'library', 'music_school', 'nightclub', 'parking', 'pharmacy', 'police', 'post_box', 'post_office', 'pub', 'public_bookcase', 'recycling', 'restaurant', 'school', 'shelter', 'social_centre', 'social_facility', 'studio', 'theatre', 'toilets', 'townhall', 'university', 'vehicle_inspection', 'veterinary']
+
+
+DICT_9to74 = {'sustenance':['bar','cafe','fast_food','food_court','pub','restaurant'],
+    'education':['college','driving_school','kindergarten','language_school','library','music_school','school','university'],
+    'transportation':['bicycle_parking','bicycle_repair_station','bicycle_rental','boat_rental',
+    'bus_station','car_rental','car_sharing','car_wash','vehicle_inspection','charging_station','ferry_terminal',
+        'fuel','motorcycle_parking','parking','taxi'],
+    'financial':['atm','bank','bureau_de_change'],
+    'healthcare':['clinic','dentist','doctors','hospital','nursing_home','pharmacy','social_facility','veterinary'],
+    'entertainment_arts_culture':['arts_centre','cinema','community_centre',
+        'conference_centre','events_venue','fountain','gambling',
+        'nightclub','public_bookcase','social_centre','stripclub','studio','swingerclub','theatre'],
+    'public_service':['courthouse','fire_station','police','post_box',
+        'post_depot','post_office','townhall'],
+    'facilities':['bench','drinking_water','parcel_locker','shelter',
+'telephone','toilets','water_point','watering_place'],
+    'waste_management':['sanitary_dump_station','recycling','waste_basket','waste_disposal','waste_transfer_station',]
+                 }
+
+DICT_74to9 = revert_dict(DICT_9to74)
+
+DICT_9to95 = {
+    'education':{'college','driving_school','kindergarten','language_school','library','toy_library','training','music_school','school','university'},
+    'entertainment_arts_culture':{'arts_centre','brothel','casino','cinema','community_centre','conference_centre','events_venue','fountain','gambling','love_hotel','nightclub','planetarium','public_bookcase','social_centre','stripclub','studio','swingerclub','theatre'},
+    'facilities':{'bbq','bench','dog_toilet','dressing_room','drinking_water','give_box','parcel_locker','shelter','shower','telephone','toilets','water_point','watering_place'},
+    'financial':{'atm','bank','bureau_de_change'},
+    'healthcare':{'baby_hatch','clinic','dentist','doctors','hospital','nursing_home','pharmacy','social_facility','veterinary'},
+    'public_service':{'courthouse','fire_station','police','post_box','post_depot','post_office','prison','ranger_station','townhall'},
+    'sustenance':{'bar','biergarten','cafe','fast_food','food_court','ice_cream','pub','restaurant',},
+    'transportation':{'bicycle_parking','bicycle_repair_station','bicycle_rental','boat_rental','boat_sharing','bus_station','car_rental','car_sharing','car_wash','compressed_air','vehicle_inspection','charging_station','ferry_terminal','fuel','grit_bin','motorcycle_parking','parking','parking_entrance','parking_space','taxi'},
+    'waste_management':{'sanitary_dump_station','recycling','waste_basket','waste_disposal','waste_transfer_station'}}
+
+DICT_95to9 = revert_dict(DICT_9to95)
+
+# DICT_95to9 = {
+#     'college':'education',
+#     'driving_school':'education',
+#     'kindergarten':'education',
+#     'language_school':'education',
+#     'library':'education',
+#     'toy_library':'education',
+#     'training':'education',
+#     'music_school':'education',
+#     'school':'education',
+#     'university':'education',
+
+#     'arts_centre':'entertainment_arts_culture',
+#     'brothel':'entertainment_arts_culture',
+#     'casino':'entertainment_arts_culture',
+#     'cinema':'entertainment_arts_culture',
+#     'community_centre':'entertainment_arts_culture',
+#     'conference_centre':'entertainment_arts_culture',
+#     'events_venue':'entertainment_arts_culture',
+#     'fountain':'entertainment_arts_culture',
+#     'gambling':'entertainment_arts_culture',
+#     'love_hotel':'entertainment_arts_culture',
+#     'nightclub':'entertainment_arts_culture',
+#     'planetarium':'entertainment_arts_culture',
+#     'public_bookcase':'entertainment_arts_culture',
+#     'social_centre':'entertainment_arts_culture',
+#     'stripclub':'entertainment_arts_culture',
+#     'studio':'entertainment_arts_culture',
+#     'swingerclub':'entertainment_arts_culture',
+#     'theatre':'entertainment_arts_culture',
+
+#     'bbq': 'facilities',
+#     'bench': 'facilities',
+#     'dog_toilet': 'facilities',
+#     'dressing_room': 'facilities',
+#     'drinking_water': 'facilities',
+#     'give_box': 'facilities',
+#     'parcel_locker': 'facilities',
+#     'shelter': 'facilities',
+#     'shower': 'facilities',
+#     'telephone': 'facilities',
+#     'toilets': 'facilities',
+#     'water_point': 'facilities',
+#     'watering_place': 'facilities',
+
+#     'atm': 'financial',
+#     'bank': 'financial',
+#     'bureau_de_change': 'financial',
+
+#     'baby_hatch':'healthcare',
+#     'clinic':'healthcare',
+#     'dentist':'healthcare',
+#     'doctors':'healthcare',
+#     'hospital':'healthcare',
+#     'nursing_home':'healthcare',
+#     'pharmacy':'healthcare',
+#     'social_facility':'healthcare',
+#     'veterinary':'healthcare',
+
+#     'courthouse': 'public_service',
+#     'fire_station': 'public_service',
+#     'police': 'public_service',
+#     'post_box': 'public_service',
+#     'post_depot': 'public_service',
+#     'post_office': 'public_service',
+#     'prison': 'public_service',
+#     'ranger_station': 'public_service',
+#     'townhall': 'public_service',
+
+#     'bar': 'sustenance',
+#     'biergarten': 'sustenance',
+#     'cafe': 'sustenance',
+#     'fast_food': 'sustenance',
+#     'food_court': 'sustenance',
+#     'ice_cream': 'sustenance',
+#     'pub': 'sustenance',
+#     'restaurant': 'sustenance',
+
+#     'bicycle_parking': 'transportation',
+#     'bicycle_repair_station': 'transportation',
+#     'bicycle_rental': 'transportation',
+#     'boat_rental': 'transportation',
+#     'boat_sharing': 'transportation',
+#     'bus_station': 'transportation',
+#     'car_rental': 'transportation',
+#     'car_sharing': 'transportation',
+#     'car_wash': 'transportation',
+#     'compressed_air': 'transportation',
+#     'vehicle_inspection': 'transportation',
+#     'charging_station': 'transportation',
+#     'ferry_terminal': 'transportation',
+#     'fuel': 'transportation',
+#     'grit_bin': 'transportation',
+#     'motorcycle_parking': 'transportation',
+#     'parking': 'transportation',
+#     'parking_entrance': 'transportation',
+#     'parking_space': 'transportation',
+#     'taxi': 'transportation',
+
+#     'sanitary_dump_station': 'waste_management',
+#     'recycling': 'waste_management',
+#     'waste_basket': 'waste_management',
+#     'waste_disposal': 'waste_management',
+#     'waste_transfer_station': 'waste_management',
+    
+# }
+
+# CLASS_9_LIST = ['sustenance', 'education', 'transportation', 'financial', 'healthcare', 'entertainment_arts_culture', 'public_service', 'facilities', 'waste_management']
+
+# FINE_LIST = ['bar','biergarten','cafe','fast_food','food_court','ice_cream','pub','restaurant','college','driving_school','kindergarten','language_school','library','toy_library','training','music_school','school','university','bicycle_parking','bicycle_repair_station','bicycle_rental','boat_rental','boat_sharing','bus_station','car_rental','car_sharing','car_wash','compressed_air','vehicle_inspection','charging_station','ferry_terminal','fuel','grit_bin','motorcycle_parking','parking','parking_entrance','parking_space','taxi','atm','bank','bureau_de_change','baby_hatch','clinic','dentist','doctors','hospital','nursing_home','pharmacy','social_facility','veterinary','arts_centre','brothel','casino','cinema','community_centre','conference_centre','events_venue','fountain','gambling','love_hotel','nightclub','planetarium','public_bookcase','social_centre','stripclub','studio','swingerclub','theatre','courthouse','fire_station','police','post_box','post_depot','post_office','prison','ranger_station','townhall','bbq','bench','dog_toilet','dressing_room','drinking_water','give_box','parcel_locker','shelter','shower','telephone','toilets','water_point','watering_place','sanitary_dump_station','recycling','waste_basket','waste_disposal','waste_transfer_station']
+
+# FINE_LIST = ['bar','biergarten','cafe','fast_food','food_court','ice_cream','pub','restaurant','college','driving_school','kindergarten','language_school','library','toy_library','training','music_school','school','university','bicycle_parking','bicycle_repair_station','bicycle_rental','boat_rental','boat_sharing','bus_station','car_rental','car_sharing','car_wash','compressed_air','vehicle_inspection','charging_station','ferry_terminal','fuel','grit_bin','motorcycle_parking','parking','parking_entrance','parking_space','taxi','atm','bank','bureau_de_change','baby_hatch','clinic','dentist','doctors','hospital','nursing_home','pharmacy','social_facility','veterinary','arts_centre','brothel','casino','cinema','community_centre','conference_centre','events_venue','fountain','gambling','love_hotel','nightclub','planetarium','public_bookcase','social_centre','stripclub','studio','swingerclub','theatre','courthouse','fire_station','police','post_box','post_depot','post_office','prison','ranger_station','townhall','bbq','bench','dog_toilet','dressing_room','drinking_water','give_box','parcel_locker','shelter','shower','telephone','toilets','water_point','watering_place','sanitary_dump_station','recycling','waste_basket','waste_disposal','waste_transfer_station','animal_boarding','animal_breeding','animal_shelter','baking_oven','childcare','clock','crematorium','dive_centre','funeral_hall','grave_yard','hunting_stand','internet_cafe','kitchen','kneipp_water_cure','lounger','marketplace','monastery','photo_booth','place_of_mourning','place_of_worship','public_bath','refugee_site','vending_machine']

models/spabert/datasets/dataset_loader.py

@@ -0,0 +1,162 @@
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+import pdb
+
+class SpatialDataset(Dataset):
+    def __init__(self, tokenizer , max_token_len ,  distance_norm_factor, sep_between_neighbors = False ):
+        self.tokenizer = tokenizer
+        self.max_token_len = max_token_len 
+        self.distance_norm_factor = distance_norm_factor
+        self.sep_between_neighbors = sep_between_neighbors
+        
+
+    def parse_spatial_context(self, pivot_name, pivot_pos, neighbor_name_list, neighbor_geometry_list, spatial_dist_fill,  pivot_dist_fill = 0):
+
+        sep_token_id = self.tokenizer.convert_tokens_to_ids(self.tokenizer.sep_token)
+        cls_token_id = self.tokenizer.convert_tokens_to_ids(self.tokenizer.cls_token)
+        mask_token_id  = self.tokenizer.convert_tokens_to_ids(self.tokenizer.mask_token)
+        pad_token_id = self.tokenizer.convert_tokens_to_ids(self.tokenizer.pad_token)
+        max_token_len = self.max_token_len
+
+
+        # process pivot
+        pivot_name_tokens = self.tokenizer.convert_tokens_to_ids(self.tokenizer.tokenize(pivot_name))
+        pivot_token_len = len(pivot_name_tokens)
+            
+        pivot_lng = pivot_pos[0]
+        pivot_lat = pivot_pos[1]
+
+        # prepare entity mask
+        entity_mask_arr = []
+        rand_entity = np.random.uniform(size = len(neighbor_name_list) + 1) # random number for masking entities including neighbors and pivot
+        # True for mask, False for unmask
+        
+        # check if pivot entity needs to be masked out, 15% prob. to be masked out
+        if rand_entity[0] < 0.15:
+            entity_mask_arr.extend([True] * pivot_token_len)
+        else:
+            entity_mask_arr.extend([False] * pivot_token_len)
+
+        # process neighbors
+        neighbor_token_list = []
+        neighbor_lng_list = []
+        neighbor_lat_list = []
+
+        # add separator between pivot and neighbor tokens
+        # a trick to avoid adding separator token after the class name (for class name encoding of margin-ranking loss)
+        if self.sep_between_neighbors and pivot_dist_fill==0: 
+            neighbor_lng_list.append(spatial_dist_fill)
+            neighbor_lat_list.append(spatial_dist_fill)
+            neighbor_token_list.append(sep_token_id)
+
+        for neighbor_name, neighbor_geometry, rnd in zip(neighbor_name_list, neighbor_geometry_list, rand_entity[1:]):
+
+            if not neighbor_name[0].isalpha():
+                # only consider neighbors starting with letters
+                continue 
+
+            neighbor_token = self.tokenizer.convert_tokens_to_ids(self.tokenizer.tokenize(neighbor_name))
+            neighbor_token_len = len(neighbor_token)
+
+            # compute the relative distance from neighbor to pivot,
+            # normalize the relative distance by distance_norm_factor
+            # apply the calculated distance for all the subtokens of the neighbor
+            # neighbor_lng_list.extend([(neighbor_geometry[0]- pivot_lng)/self.distance_norm_factor] * neighbor_token_len)
+            # neighbor_lat_list.extend([(neighbor_geometry[1]- pivot_lat)/self.distance_norm_factor] * neighbor_token_len)
+
+            if 'coordinates' in neighbor_geometry: # to handle different json dict structures
+                neighbor_lng_list.extend([(neighbor_geometry['coordinates'][0]- pivot_lng)/self.distance_norm_factor] * neighbor_token_len)
+                neighbor_lat_list.extend([(neighbor_geometry['coordinates'][1]- pivot_lat)/self.distance_norm_factor] * neighbor_token_len)
+                neighbor_token_list.extend(neighbor_token)
+            else:
+                neighbor_lng_list.extend([(neighbor_geometry[0]- pivot_lng)/self.distance_norm_factor] * neighbor_token_len)
+                neighbor_lat_list.extend([(neighbor_geometry[1]- pivot_lat)/self.distance_norm_factor] * neighbor_token_len)
+                neighbor_token_list.extend(neighbor_token)
+
+            if self.sep_between_neighbors:
+                neighbor_lng_list.append(spatial_dist_fill)
+                neighbor_lat_list.append(spatial_dist_fill)
+                neighbor_token_list.append(sep_token_id)
+                
+                entity_mask_arr.extend([False])
+
+            
+            if rnd < 0.15:
+                #True: mask out, False: Keey original token
+                entity_mask_arr.extend([True] * neighbor_token_len)
+            else:
+                entity_mask_arr.extend([False] * neighbor_token_len)
+
+
+        pseudo_sentence = pivot_name_tokens + neighbor_token_list 
+        dist_lng_list = [pivot_dist_fill] * pivot_token_len + neighbor_lng_list 
+        dist_lat_list = [pivot_dist_fill] * pivot_token_len + neighbor_lat_list 
+        
+
+        #including cls and sep
+        sent_len = len(pseudo_sentence)
+
+        max_token_len_middle = max_token_len -2 # 2 for CLS and SEP token
+
+        # padding and truncation
+        if sent_len > max_token_len_middle : 
+            pseudo_sentence = [cls_token_id] + pseudo_sentence[:max_token_len_middle] + [sep_token_id] 
+            dist_lat_list = [spatial_dist_fill] + dist_lat_list[:max_token_len_middle]+ [spatial_dist_fill]
+            dist_lng_list = [spatial_dist_fill] + dist_lng_list[:max_token_len_middle]+ [spatial_dist_fill]
+            attention_mask = [False] + [1] * max_token_len_middle + [False] # make sure SEP and CLS are not attented to
+        else:
+            pad_len = max_token_len_middle - sent_len
+            assert pad_len >= 0 
+
+            pseudo_sentence = [cls_token_id] + pseudo_sentence + [sep_token_id] + [pad_token_id] * pad_len 
+            dist_lat_list = [spatial_dist_fill] + dist_lat_list + [spatial_dist_fill] + [spatial_dist_fill] * pad_len
+            dist_lng_list = [spatial_dist_fill] + dist_lng_list + [spatial_dist_fill] + [spatial_dist_fill] * pad_len
+            attention_mask = [False] + [1] * sent_len + [0] * pad_len + [False]
+
+
+        
+
+        norm_lng_list = np.array(dist_lng_list) # / 0.0001
+        norm_lat_list = np.array(dist_lat_list) # / 0.0001
+
+
+        # mask entity in the pseudo sentence 
+        entity_mask_indices = np.where(entity_mask_arr)[0]
+        masked_entity_input = [mask_token_id if i in entity_mask_indices else pseudo_sentence[i] for i in range(0, max_token_len)]
+
+        
+        # mask token in the pseudo sentence
+        rand_token = np.random.uniform(size = len(pseudo_sentence))
+        # do not mask out cls and sep token. True: masked tokens False: Keey original token
+        token_mask_arr = (rand_token <0.15) & (np.array(pseudo_sentence) != cls_token_id) & (np.array(pseudo_sentence) != sep_token_id) & (np.array(pseudo_sentence) != pad_token_id)
+        token_mask_indices = np.where(token_mask_arr)[0]
+        
+        masked_token_input = [mask_token_id if i in token_mask_indices else pseudo_sentence[i] for i in range(0, max_token_len)]
+        
+
+        # yield masked_token with 50% prob, masked_entity with 50% prob
+        if np.random.rand() > 0.5:
+            masked_input = torch.tensor(masked_entity_input)
+        else:
+            masked_input = torch.tensor(masked_token_input)
+        
+        train_data = {}
+        train_data['pivot_name'] = pivot_name
+        train_data['pivot_token_len'] = pivot_token_len
+        train_data['masked_input'] = masked_input
+        train_data['sent_position_ids'] = torch.tensor(np.arange(0, len(pseudo_sentence)))
+        train_data['attention_mask'] = torch.tensor(attention_mask)
+        train_data['norm_lng_list'] = torch.tensor(norm_lng_list).to(torch.float32)
+        train_data['norm_lat_list'] = torch.tensor(norm_lat_list).to(torch.float32)
+        train_data['pseudo_sentence'] = torch.tensor(pseudo_sentence)
+
+        return train_data
+
+
+
+    def __len__(self):
+        return NotImplementedError
+
+    def __getitem__(self, index):
+        raise NotImplementedError

models/spabert/datasets/dataset_loader_ver2.py

@@ -0,0 +1,164 @@
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+import pdb
+
+class SpatialDataset(Dataset):
+    def __init__(self, tokenizer , max_token_len ,  distance_norm_factor, sep_between_neighbors = False ):
+        self.tokenizer = tokenizer
+        self.max_token_len = max_token_len 
+        self.distance_norm_factor = distance_norm_factor
+        self.sep_between_neighbors = sep_between_neighbors
+        
+
+    def parse_spatial_context(self, pivot_name, pivot_pos, neighbor_name_list, neighbor_geometry_list, spatial_dist_fill,  pivot_dist_fill = 0):
+
+        sep_token_id = self.tokenizer.convert_tokens_to_ids(self.tokenizer.sep_token)
+        cls_token_id = self.tokenizer.convert_tokens_to_ids(self.tokenizer.cls_token)
+        #mask_token_id  = self.tokenizer.convert_tokens_to_ids(self.tokenizer.mask_token)
+        pad_token_id = self.tokenizer.convert_tokens_to_ids(self.tokenizer.pad_token)
+        max_token_len = self.max_token_len
+
+       
+        #print("Module reloaded and changes are reflected")
+        # process pivot
+        pivot_name_tokens = self.tokenizer.convert_tokens_to_ids(self.tokenizer.tokenize(pivot_name))
+        pivot_token_len = len(pivot_name_tokens)
+            
+        pivot_lng = pivot_pos[0]
+        pivot_lat = pivot_pos[1]
+
+        # prepare entity mask
+        entity_mask_arr = []
+        rand_entity = np.random.uniform(size = len(neighbor_name_list) + 1) # random number for masking entities including neighbors and pivot
+        # True for mask, False for unmask
+        
+        # check if pivot entity needs to be masked out, 15% prob. to be masked out
+        #if rand_entity[0] < 0.15:
+        #    entity_mask_arr.extend([True] * pivot_token_len)
+        #else:
+        entity_mask_arr.extend([False] * pivot_token_len)
+
+        # process neighbors
+        neighbor_token_list = []
+        neighbor_lng_list = []
+        neighbor_lat_list = []
+
+        # add separator between pivot and neighbor tokens
+        # a trick to avoid adding separator token after the class name (for class name encoding of margin-ranking loss)
+        if self.sep_between_neighbors and pivot_dist_fill==0: 
+            neighbor_lng_list.append(spatial_dist_fill)
+            neighbor_lat_list.append(spatial_dist_fill)
+            neighbor_token_list.append(sep_token_id)
+
+        for neighbor_name, neighbor_geometry, rnd in zip(neighbor_name_list, neighbor_geometry_list, rand_entity[1:]):
+
+            if not neighbor_name[0].isalpha():
+                # only consider neighbors starting with letters
+                continue 
+
+            neighbor_token = self.tokenizer.convert_tokens_to_ids(self.tokenizer.tokenize(neighbor_name))
+            neighbor_token_len = len(neighbor_token)
+
+            # compute the relative distance from neighbor to pivot,
+            # normalize the relative distance by distance_norm_factor
+            # apply the calculated distance for all the subtokens of the neighbor
+            # neighbor_lng_list.extend([(neighbor_geometry[0]- pivot_lng)/self.distance_norm_factor] * neighbor_token_len)
+            # neighbor_lat_list.extend([(neighbor_geometry[1]- pivot_lat)/self.distance_norm_factor] * neighbor_token_len)
+
+            if 'coordinates' in neighbor_geometry: # to handle different json dict structures
+                neighbor_lng_list.extend([(neighbor_geometry['coordinates'][0]- pivot_lng)/self.distance_norm_factor] * neighbor_token_len)
+                neighbor_lat_list.extend([(neighbor_geometry['coordinates'][1]- pivot_lat)/self.distance_norm_factor] * neighbor_token_len)
+                neighbor_token_list.extend(neighbor_token)
+            else:
+                neighbor_lng_list.extend([(neighbor_geometry[0]- pivot_lng)/self.distance_norm_factor] * neighbor_token_len)
+                neighbor_lat_list.extend([(neighbor_geometry[1]- pivot_lat)/self.distance_norm_factor] * neighbor_token_len)
+                neighbor_token_list.extend(neighbor_token)
+
+            if self.sep_between_neighbors:
+                neighbor_lng_list.append(spatial_dist_fill)
+                neighbor_lat_list.append(spatial_dist_fill)
+                neighbor_token_list.append(sep_token_id)
+                
+                entity_mask_arr.extend([False])
+
+            
+            #if rnd < 0.15:
+            #    #True: mask out, False: Keey original token
+            #    entity_mask_arr.extend([True] * neighbor_token_len)
+            #else:
+            entity_mask_arr.extend([False] * neighbor_token_len)
+
+
+        pseudo_sentence = pivot_name_tokens + neighbor_token_list 
+        dist_lng_list = [pivot_dist_fill] * pivot_token_len + neighbor_lng_list 
+        dist_lat_list = [pivot_dist_fill] * pivot_token_len + neighbor_lat_list 
+        
+
+        #including cls and sep
+        sent_len = len(pseudo_sentence)
+
+        max_token_len_middle = max_token_len -2 # 2 for CLS and SEP token
+
+        # padding and truncation
+        if sent_len > max_token_len_middle : 
+            pseudo_sentence = [cls_token_id] + pseudo_sentence[:max_token_len_middle] + [sep_token_id] 
+            dist_lat_list = [spatial_dist_fill] + dist_lat_list[:max_token_len_middle]+ [spatial_dist_fill]
+            dist_lng_list = [spatial_dist_fill] + dist_lng_list[:max_token_len_middle]+ [spatial_dist_fill]
+            attention_mask = [False] + [1] * max_token_len_middle + [False] # make sure SEP and CLS are not attented to
+        else:
+            pad_len = max_token_len_middle - sent_len
+            assert pad_len >= 0 
+
+            pseudo_sentence = [cls_token_id] + pseudo_sentence + [sep_token_id] + [pad_token_id] * pad_len 
+            dist_lat_list = [spatial_dist_fill] + dist_lat_list + [spatial_dist_fill] + [spatial_dist_fill] * pad_len
+            dist_lng_list = [spatial_dist_fill] + dist_lng_list + [spatial_dist_fill] + [spatial_dist_fill] * pad_len
+            attention_mask = [False] + [1] * sent_len + [0] * pad_len + [False]
+
+
+        
+
+        norm_lng_list = np.array(dist_lng_list) # / 0.0001
+        norm_lat_list = np.array(dist_lat_list) # / 0.0001
+
+
+        ## mask entity in the pseudo sentence 
+        #entity_mask_indices = np.where(entity_mask_arr)[0]
+        #masked_entity_input = [mask_token_id if i in entity_mask_indices else pseudo_sentence[i] for i in range(0, max_token_len)]
+        #
+        #
+        ## mask token in the pseudo sentence
+        #rand_token = np.random.uniform(size = len(pseudo_sentence))
+        ## do not mask out cls and sep token. True: masked tokens False: Keey original token
+        #token_mask_arr = (rand_token <0.15) & (np.array(pseudo_sentence) != cls_token_id) & (np.array(pseudo_sentence) != sep_token_id) & (np.array(pseudo_sentence) != pad_token_id)
+        #token_mask_indices = np.where(token_mask_arr)[0]
+        #
+        #masked_token_input = [mask_token_id if i in token_mask_indices else pseudo_sentence[i] for i in range(0, max_token_len)]
+        #
+        #
+        ## yield masked_token with 50% prob, masked_entity with 50% prob
+        #if np.random.rand() > 0.5:
+        #    masked_input = torch.tensor(masked_entity_input)
+        #else:
+        #    masked_input = torch.tensor(masked_token_input)
+        masked_input = torch.tensor(pseudo_sentence)
+        
+        train_data = {}
+        train_data['pivot_name'] = pivot_name
+        train_data['pivot_token_len'] = pivot_token_len
+        train_data['masked_input'] = masked_input
+        train_data['sent_position_ids'] = torch.tensor(np.arange(0, len(pseudo_sentence)))
+        train_data['attention_mask'] = torch.tensor(attention_mask)
+        train_data['norm_lng_list'] = torch.tensor(norm_lng_list).to(torch.float32)
+        train_data['norm_lat_list'] = torch.tensor(norm_lat_list).to(torch.float32)
+        train_data['pseudo_sentence'] = torch.tensor(pseudo_sentence)
+
+        return train_data
+
+
+
+    def __len__(self):
+        return NotImplementedError
+
+    def __getitem__(self, index):
+        raise NotImplementedError

models/spabert/datasets/osm_sample_loader.py

@@ -0,0 +1,246 @@
+import os
+import sys
+import numpy as np
+import json 
+import math
+
+import torch
+from transformers import RobertaTokenizer, BertTokenizer
+from torch.utils.data import Dataset
+#sys.path.append('/home/zekun/spatial_bert/spatial_bert/datasets')
+sys.path.append('/content/drive/MyDrive/spaBERT/spabert/datasets')
+from dataset_loader_ver2 import SpatialDataset
+
+import pdb
+
+class PbfMapDataset(SpatialDataset):
+    def __init__(self, data_file_path,  tokenizer=None, max_token_len = 512, distance_norm_factor = 0.0001, spatial_dist_fill=10, 
+        with_type = True, sep_between_neighbors = False, label_encoder = None, mode = None, num_neighbor_limit = None, random_remove_neighbor = 0.,type_key_str='class'):
+        
+        if tokenizer is None:
+            self.tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+        else:
+            self.tokenizer = tokenizer
+
+        self.max_token_len = max_token_len
+        self.spatial_dist_fill = spatial_dist_fill # should be normalized distance fill, larger than all normalized neighbor distance
+        self.with_type = with_type
+        self.sep_between_neighbors = sep_between_neighbors
+        self.label_encoder = label_encoder
+        self.num_neighbor_limit = num_neighbor_limit
+        self.read_file(data_file_path, mode)
+        self.random_remove_neighbor = random_remove_neighbor
+        self.type_key_str  = type_key_str # key name of the class type in the input data dictionary
+ 
+        super(PbfMapDataset, self).__init__(self.tokenizer , max_token_len , distance_norm_factor, sep_between_neighbors )
+        
+
+    def read_file(self, data_file_path, mode):
+
+        with open(data_file_path, 'r') as f:
+            data = f.readlines()
+
+        if mode == 'train':
+            data = data[0:int(len(data) * 0.8)]
+        elif mode == 'test':
+            data = data[int(len(data) * 0.8):]
+        elif mode is None: # use the full dataset (for mlm)
+            pass
+        else:
+            raise NotImplementedError
+
+        self.len_data = len(data) # updated data length
+        self.data = data 
+
+    def load_data(self, index):
+        
+        spatial_dist_fill = self.spatial_dist_fill
+        line = self.data[index] # take one line from the input data according to the index
+
+        line_data_dict = json.loads(line)
+
+        # process pivot
+        pivot_name = line_data_dict['info']['name']
+        pivot_pos = line_data_dict['info']['geometry']['coordinates']
+
+        
+        neighbor_info = line_data_dict['neighbor_info']
+        neighbor_name_list = neighbor_info['name_list']
+        neighbor_geometry_list = neighbor_info['geometry_list']
+
+        if self.random_remove_neighbor != 0:
+            num_neighbors = len(neighbor_name_list)
+            rand_neighbor = np.random.uniform(size = num_neighbors)
+
+            neighbor_keep_arr = (rand_neighbor >= self.random_remove_neighbor) # select the neighbors to be removed
+            neighbor_keep_arr = np.where(neighbor_keep_arr)[0]
+            
+            new_neighbor_name_list, new_neighbor_geometry_list = [],[]
+            for i in range(0, num_neighbors):
+                if i in neighbor_keep_arr:
+                    new_neighbor_name_list.append(neighbor_name_list[i])
+                    new_neighbor_geometry_list.append(neighbor_geometry_list[i])
+
+            neighbor_name_list = new_neighbor_name_list
+            neighbor_geometry_list = new_neighbor_geometry_list
+        
+        if self.num_neighbor_limit is not None:
+            neighbor_name_list = neighbor_name_list[0:self.num_neighbor_limit]
+            neighbor_geometry_list = neighbor_geometry_list[0:self.num_neighbor_limit]
+
+
+        train_data = self.parse_spatial_context(pivot_name, pivot_pos, neighbor_name_list, neighbor_geometry_list, spatial_dist_fill )
+
+        if self.with_type:
+            pivot_type = line_data_dict['info'][self.type_key_str]
+            train_data['pivot_type'] = torch.tensor(self.label_encoder.transform([pivot_type])[0]) # scalar, label_id
+
+        if 'ogc_fid' in line_data_dict['info']:
+            train_data['ogc_fid'] = line_data_dict['info']['ogc_fid']
+
+        return train_data
+
+    def __len__(self):
+        return self.len_data
+
+    def __getitem__(self, index):
+        return self.load_data(index)
+
+
+
+class PbfMapDatasetMarginRanking(SpatialDataset):
+    def __init__(self, data_file_path,  type_list = None, tokenizer=None, max_token_len = 512, distance_norm_factor = 0.0001, spatial_dist_fill=10, 
+        sep_between_neighbors = False,  mode = None, num_neighbor_limit = None, random_remove_neighbor = 0., type_key_str='class'):
+        
+        if tokenizer is None:
+            self.tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+        else:
+            self.tokenizer = tokenizer
+
+        self.type_list  = type_list
+        self.type_key_str  = type_key_str # key name of the class type in the input data dictionary
+        self.max_token_len = max_token_len
+        self.spatial_dist_fill = spatial_dist_fill # should be normalized distance fill, larger than all normalized neighbor distance
+        self.sep_between_neighbors = sep_between_neighbors
+        # self.label_encoder = label_encoder
+        self.num_neighbor_limit = num_neighbor_limit
+        self.read_file(data_file_path, mode)
+        self.random_remove_neighbor = random_remove_neighbor
+        self.mode = mode
+        
+ 
+        super(PbfMapDatasetMarginRanking, self).__init__(self.tokenizer , max_token_len , distance_norm_factor, sep_between_neighbors )
+        
+
+    def read_file(self, data_file_path, mode):
+
+        with open(data_file_path, 'r') as f:
+            data = f.readlines()
+
+        if mode == 'train':
+            data = data[0:int(len(data) * 0.8)]
+        elif mode == 'test':
+            data = data[int(len(data) * 0.8):]
+            self.all_types_data = self.prepare_all_types_data()
+        elif mode is None: # use the full dataset (for mlm)
+            pass
+        else:
+            raise NotImplementedError
+
+        self.len_data = len(data) # updated data length
+        self.data = data 
+
+    def prepare_all_types_data(self):
+        type_list = self.type_list
+        spatial_dist_fill = self.spatial_dist_fill
+        type_data_dict = dict()
+        for type_name in type_list:
+            type_pos = [None, None] # use filler values
+            type_data = self.parse_spatial_context(type_name, type_pos, pivot_dist_fill = 0.,
+                neighbor_name_list = [], neighbor_geometry_list=[], spatial_dist_fill= spatial_dist_fill)
+            type_data_dict[type_name] = type_data
+
+        return type_data_dict
+
+    def load_data(self, index):
+        
+        spatial_dist_fill = self.spatial_dist_fill
+        line = self.data[index] # take one line from the input data according to the index
+
+        line_data_dict = json.loads(line)
+
+        # process pivot
+        pivot_name = line_data_dict['info']['name']
+        pivot_pos = line_data_dict['info']['geometry']['coordinates']
+
+        
+        neighbor_info = line_data_dict['neighbor_info']
+        neighbor_name_list = neighbor_info['name_list']
+        neighbor_geometry_list = neighbor_info['geometry_list']
+
+        if self.random_remove_neighbor != 0:
+            num_neighbors = len(neighbor_name_list)
+            rand_neighbor = np.random.uniform(size = num_neighbors)
+
+            neighbor_keep_arr = (rand_neighbor >= self.random_remove_neighbor) # select the neighbors to be removed
+            neighbor_keep_arr = np.where(neighbor_keep_arr)[0]
+            
+            new_neighbor_name_list, new_neighbor_geometry_list = [],[]
+            for i in range(0, num_neighbors):
+                if i in neighbor_keep_arr:
+                    new_neighbor_name_list.append(neighbor_name_list[i])
+                    new_neighbor_geometry_list.append(neighbor_geometry_list[i])
+
+            neighbor_name_list = new_neighbor_name_list
+            neighbor_geometry_list = new_neighbor_geometry_list
+        
+        if self.num_neighbor_limit is not None:
+            neighbor_name_list = neighbor_name_list[0:self.num_neighbor_limit]
+            neighbor_geometry_list = neighbor_geometry_list[0:self.num_neighbor_limit]
+
+
+        train_data = self.parse_spatial_context(pivot_name, pivot_pos, neighbor_name_list, neighbor_geometry_list, spatial_dist_fill )
+
+        if 'ogc_fid' in line_data_dict['info']:
+            train_data['ogc_fid'] = line_data_dict['info']['ogc_fid']
+
+        # train_data['pivot_type'] = torch.tensor(self.label_encoder.transform([pivot_type])[0]) # scalar, label_id
+        
+        pivot_type = line_data_dict['info'][self.type_key_str]
+        train_data['pivot_type'] = pivot_type
+
+        if self.mode == 'train':
+            # postive class
+            postive_name = pivot_type # class type string as input to tokenizer
+            positive_pos = [None, None] # use filler values
+            postive_type_data = self.parse_spatial_context(postive_name, positive_pos, pivot_dist_fill = 0.,
+                neighbor_name_list = [], neighbor_geometry_list=[], spatial_dist_fill= spatial_dist_fill)
+            train_data['positive_type_data'] = postive_type_data
+
+
+            # negative class
+            other_type_list = self.type_list.copy()
+            other_type_list.remove(pivot_type)
+            other_type = np.random.choice(other_type_list)
+            negative_name = other_type 
+            negative_pos = [None, None] # use filler values
+            negative_type_data = self.parse_spatial_context(negative_name, negative_pos, pivot_dist_fill = 0.,
+                neighbor_name_list = [], neighbor_geometry_list=[], spatial_dist_fill= spatial_dist_fill)
+            train_data['negative_type_data'] = negative_type_data
+
+        elif self.mode == 'test':
+            # return data for all class types in type_list
+            train_data['all_types_data'] = self.all_types_data
+
+        else:
+            raise NotImplementedError
+
+        return train_data
+
+    def __len__(self):
+        return self.len_data
+
+    def __getitem__(self, index):
+        return self.load_data(index)
+
+

models/spabert/datasets/usgs_os_sample_loader.py

@@ -0,0 +1,71 @@
+import os
+import sys
+import numpy as np
+import json 
+import math
+from sklearn.preprocessing import LabelBinarizer, LabelEncoder
+
+import torch
+from transformers import RobertaTokenizer, BertTokenizer
+from torch.utils.data import Dataset
+sys.path.append('/home/zekun/spatial_bert/spatial_bert/datasets')
+from dataset_loader import SpatialDataset
+
+import pdb
+
+
+class USGS_MapDataset(SpatialDataset):
+    def __init__(self, data_file_path,  tokenizer=None, max_token_len = 512, distance_norm_factor = 1, spatial_dist_fill=100, sep_between_neighbors = False):
+        
+        if tokenizer is None:
+            self.tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+        else:
+            self.tokenizer = tokenizer
+
+        self.max_token_len = max_token_len
+        self.spatial_dist_fill = spatial_dist_fill # should be normalized distance fill, larger than all normalized neighbor distance
+        self.sep_between_neighbors = sep_between_neighbors
+        self.read_file(data_file_path)
+
+        
+        super(USGS_MapDataset, self).__init__(self.tokenizer , max_token_len , distance_norm_factor, sep_between_neighbors )
+        
+    def read_file(self, data_file_path):
+
+        with open(data_file_path, 'r') as f:
+            data = f.readlines()
+
+        len_data = len(data)
+        self.len_data = len_data
+        self.data = data 
+
+
+    def load_data(self, index):
+        
+        spatial_dist_fill = self.spatial_dist_fill
+        line = self.data[index] # take one line from the input data according to the index
+
+        line_data_dict = json.loads(line)
+
+        # process pivot
+        pivot_name = line_data_dict['info']['name']
+        pivot_pos = line_data_dict['info']['geometry']
+
+        neighbor_info = line_data_dict['neighbor_info']
+        neighbor_name_list = neighbor_info['name_list']
+        neighbor_geometry_list = neighbor_info['geometry_list']
+
+        parsed_data = self.parse_spatial_context(pivot_name, pivot_pos, neighbor_name_list, neighbor_geometry_list, spatial_dist_fill )
+
+        return parsed_data
+
+
+
+    def __len__(self):
+        return self.len_data
+
+    def __getitem__(self, index):
+        return self.load_data(index)
+
+
+

models/spabert/datasets/wikidata_sample_loader.py

@@ -0,0 +1,127 @@
+import sys
+import numpy as np
+import json 
+import math
+
+import torch
+from transformers import RobertaTokenizer, BertTokenizer
+from torch.utils.data import Dataset
+sys.path.append('/home/zekun/spatial_bert/spatial_bert/datasets')
+from dataset_loader import SpatialDataset
+
+import pdb
+
+'''Prepare candiate list given randomly sampled data and append to data_list'''
+class Wikidata_Random_Dataset(SpatialDataset): 
+    def __init__(self, data_file_path,  tokenizer=None, max_token_len = 512, distance_norm_factor = 0.0001, spatial_dist_fill=100, sep_between_neighbors = False):
+        if tokenizer is None:
+            self.tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+        else:
+            self.tokenizer = tokenizer
+
+        self.max_token_len = max_token_len
+        self.spatial_dist_fill = spatial_dist_fill # should be normalized distance fill, larger than all normalized neighbor distance
+        self.sep_between_neighbors = sep_between_neighbors
+        self.read_file(data_file_path)
+
+        
+        super(Wikidata_Random_Dataset, self).__init__(self.tokenizer , max_token_len , distance_norm_factor, sep_between_neighbors )
+
+    def read_file(self, data_file_path):
+
+        with open(data_file_path, 'r') as f:
+            data = f.readlines()
+
+        len_data = len(data)
+        self.len_data = len_data
+        self.data = data 
+
+    def load_data(self, index):
+        
+        spatial_dist_fill = self.spatial_dist_fill
+        line = self.data[index] # take one line from the input data according to the index
+
+        line_data_dict = json.loads(line)
+
+        # process pivot
+        pivot_name = line_data_dict['info']['name']
+        pivot_pos = line_data_dict['info']['geometry']['coordinates']
+        pivot_uri = line_data_dict['info']['uri']
+
+        neighbor_info = line_data_dict['neighbor_info']
+        neighbor_name_list = neighbor_info['name_list']
+        neighbor_geometry_list = neighbor_info['geometry_list']
+
+        parsed_data = self.parse_spatial_context(pivot_name, pivot_pos, neighbor_name_list, neighbor_geometry_list, spatial_dist_fill )
+        parsed_data['uri'] = pivot_uri
+        parsed_data['description'] = None # placeholder
+
+        return parsed_data
+
+    def __len__(self):
+        return self.len_data
+
+    def __getitem__(self, index):
+        return self.load_data(index)
+
+
+
+'''Prepare candiate list for each phrase and append to data_list'''
+
+class Wikidata_Geocoord_Dataset(SpatialDataset): 
+
+    #DEFAULT_CONFIG_CLS = SpatialBertConfig
+    def __init__(self, data_file_path,  tokenizer=None, max_token_len = 512, distance_norm_factor = 0.0001, spatial_dist_fill=100, sep_between_neighbors = False):
+        if tokenizer is None:
+            self.tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+        else:
+            self.tokenizer = tokenizer
+
+        self.max_token_len = max_token_len
+        self.spatial_dist_fill = spatial_dist_fill # should be normalized distance fill, larger than all normalized neighbor distance
+        self.sep_between_neighbors = sep_between_neighbors
+        self.read_file(data_file_path)
+
+        
+        super(Wikidata_Geocoord_Dataset, self).__init__(self.tokenizer , max_token_len , distance_norm_factor, sep_between_neighbors )
+
+    def read_file(self, data_file_path):
+
+        with open(data_file_path, 'r') as f:
+            data = f.readlines()
+
+        len_data = len(data)
+        self.len_data = len_data
+        self.data = data 
+
+    def load_data(self, index):
+        
+        spatial_dist_fill = self.spatial_dist_fill
+        line = self.data[index] # take one line from the input data according to the index
+
+        line_data = json.loads(line)
+        parsed_data_list = []
+
+        for line_data_dict in line_data:
+            # process pivot
+            pivot_name = line_data_dict['info']['name']
+            pivot_pos = line_data_dict['info']['geometry']['coordinates']
+            pivot_uri = line_data_dict['info']['uri']
+
+            neighbor_info = line_data_dict['neighbor_info']
+            neighbor_name_list = neighbor_info['name_list']
+            neighbor_geometry_list = neighbor_info['geometry_list']
+
+            parsed_data = self.parse_spatial_context(pivot_name, pivot_pos, neighbor_name_list, neighbor_geometry_list, spatial_dist_fill )
+            parsed_data['uri'] = pivot_uri
+            parsed_data['description'] = None # placeholder
+            parsed_data_list.append(parsed_data)
+
+        return parsed_data_list
+
+
+    def __len__(self):
+        return self.len_data
+
+    def __getitem__(self, index):
+        return self.load_data(index)

models/spabert/experiments/entity_matching/data_processing/get_namelist.py

@@ -0,0 +1,95 @@
+import json
+import os
+
+def get_name_list_osm(ref_paths):
+    name_list = []
+
+    for json_path in ref_paths:
+        with open(json_path, 'r') as f:
+            data = f.readlines()
+        for line in data:
+            record = json.loads(line)
+            name = record['name']
+            name_list.append(name)
+
+    namelist = sorted(namelist)
+    return name_list
+
+# deprecated
+def get_name_list_usgs_od(ref_paths):
+    name_list = []
+
+    for json_path in ref_paths:
+        with open(json_path, 'r') as f:
+            annot_dict = json.load(f)
+        for key, place in annot_dict.items():
+            place_name = ''
+            for idx in range(1, len(place)+1):
+                try:
+                    place_name += place[str(idx)]['text_label']
+                    place_name += ' ' # separate words with spaces
+
+                except Exception as e:
+                    print(place)
+            place_name = place_name[:-1] # remove last space
+            
+            name_list.append(place_name)
+
+    namelist = sorted(namelist)
+    return name_list
+
+def get_name_list_usgs_od_per_map(ref_paths):
+    all_name_list_dict = dict()
+
+    for json_path in ref_paths:
+        map_name = os.path.basename(json_path).split('.json')[0]
+
+        with open(json_path, 'r') as f:
+            annot_dict = json.load(f)
+
+        map_name_list = []
+        for key, place in annot_dict.items():
+            place_name = ''
+            for idx in range(1, len(place)+1):
+                try:
+                    place_name += place[str(idx)]['text_label']
+                    place_name += ' ' # separate words with spaces
+
+                except Exception as e:
+                    print(place)
+            place_name = place_name[:-1] # remove last space
+            
+            map_name_list.append(place_name)
+        all_name_list_dict[map_name] = sorted(map_name_list)
+
+    return all_name_list_dict
+
+
+def get_name_list_gb1900(ref_path):
+    name_list = []
+
+    with open(ref_path, 'r',encoding='utf-16') as f:
+        data = f.readlines()
+
+
+    for line in data[1:]: # skip the header
+        try:
+            line = line.split(',')
+            text = line[1]
+            lat = float(line[-3])
+            lng = float(line[-2])
+            semantic_type = line[-1]
+            
+            name_list.append(text)
+        except:
+            print(line)
+
+    namelist = sorted(namelist)
+
+    return name_list
+
+
+if __name__ == '__main__':
+    #name_list = get_name_list_usgs_od(['labGISReport-master/output/USGS-15-CA-brawley-e1957-s1957-p1961.json', 
+        #'labGISReport-master/output/USGS-15-CA-capesanmartin-e1921-s1917.json'])
+    name_list = get_name_list_gb1900('data/GB1900_gazetteer_abridged_july_2018/gb1900_abridged.csv')

models/spabert/experiments/entity_matching/data_processing/request_wrapper.py

@@ -0,0 +1,186 @@
+import requests
+import pdb
+import time
+
+# for linkedgeodata: #http://linkedgeodata.org/sparql
+
+class RequestWrapper:
+    def __init__(self, baseuri = "https://query.wikidata.org/sparql"):
+
+        self.baseuri = baseuri
+
+    def response_handler(self, response, query):
+        if response.status_code == requests.codes.ok: 
+            ret_json = response.json()['results']['bindings']
+        elif response.status_code == 500: 
+            ret_json = []
+            #print(q_id)
+            print('Internal Error happened. Set ret_json to be empty list')
+        
+        elif response.status_code == 429: 
+
+            print(response.status_code)
+            print(response.text)
+            retry_seconds = int(response.text.split('Too Many Requests - Please retry in ')[1].split(' seconds')[0])
+            print('rerun in %d seconds' %retry_seconds)
+            time.sleep(retry_seconds + 1)
+            
+            response = requests.get(self.baseuri, params = {'format':'json', 'query':query})
+            ret_json = response.json()['results']['bindings']
+            #print(ret_json)
+            print('resumed and succeeded')
+
+        else:
+            print(response.status_code, response.text)
+            exit(-1)
+
+        return ret_json
+
+    '''Search for wikidata entities given the name string'''
+    def wikidata_query (self, name_str): 
+        
+        query = """
+            PREFIX wd: <http://www.wikidata.org/entity/>
+            PREFIX wds: <http://www.wikidata.org/entity/statement/>
+            PREFIX wdv: <http://www.wikidata.org/value/>
+            PREFIX wdt: <http://www.wikidata.org/prop/direct/>
+            PREFIX wikibase: <http://wikiba.se/ontology#>
+            PREFIX p: <http://www.wikidata.org/prop/>
+            PREFIX ps: <http://www.wikidata.org/prop/statement/>
+            PREFIX pq: <http://www.wikidata.org/prop/qualifier/>
+            PREFIX rdfs: <http://www.w3.org/2000/01/rdf-schema#>
+            PREFIX bd: <http://www.bigdata.com/rdf#>
+
+            SELECT ?item ?coordinates ?itemDescription WHERE {
+            ?item rdfs:label \"%s\"@en;
+                    wdt:P625 ?coordinates . 
+                    SERVICE wikibase:label { bd:serviceParam wikibase:language "en" }
+                    }
+            """%(name_str)
+
+        response = requests.get(self.baseuri, params = {'format':'json', 'query':query})
+
+
+        ret_json = self.response_handler(response, query)
+
+        return ret_json
+
+
+    '''Search for wikidata entities given the name string'''
+    def wikidata_query_withinstate (self, name_str, state_id = 'Q99'): 
+
+
+        query = """
+            PREFIX wd: <http://www.wikidata.org/entity/>
+            PREFIX wds: <http://www.wikidata.org/entity/statement/>
+            PREFIX wdv: <http://www.wikidata.org/value/>
+            PREFIX wdt: <http://www.wikidata.org/prop/direct/>
+            PREFIX wikibase: <http://wikiba.se/ontology#>
+            PREFIX p: <http://www.wikidata.org/prop/>
+            PREFIX ps: <http://www.wikidata.org/prop/statement/>
+            PREFIX pq: <http://www.wikidata.org/prop/qualifier/>
+            PREFIX rdfs: <http://www.w3.org/2000/01/rdf-schema#>
+            PREFIX bd: <http://www.bigdata.com/rdf#>
+
+            SELECT ?item ?coordinates ?itemDescription WHERE {
+            ?item rdfs:label \"%s\"@en;
+                    wdt:P625 ?coordinates ;
+                    wdt:P131+ wd:%s;
+                    SERVICE wikibase:label { bd:serviceParam wikibase:language "en" }
+                    }
+            """%(name_str, state_id)
+
+        #print(query)
+
+        response = requests.get(self.baseuri, params = {'format':'json', 'query':query})
+
+        ret_json = self.response_handler(response, query)
+        
+        return ret_json
+        
+
+    '''Search for nearby wikidata entities given the entity id'''
+    def wikidata_nearby_query (self, q_id):
+
+        query = """
+            PREFIX wd: <http://www.wikidata.org/entity/>
+            PREFIX wds: <http://www.wikidata.org/entity/statement/>
+            PREFIX wdv: <http://www.wikidata.org/value/>
+            PREFIX wdt: <http://www.wikidata.org/prop/direct/>
+            PREFIX wikibase: <http://wikiba.se/ontology#>
+            PREFIX p: <http://www.wikidata.org/prop/>
+            PREFIX ps: <http://www.wikidata.org/prop/statement/>
+            PREFIX pq: <http://www.wikidata.org/prop/qualifier/>
+            PREFIX rdfs: <http://www.w3.org/2000/01/rdf-schema#>
+            PREFIX bd: <http://www.bigdata.com/rdf#>
+
+            SELECT ?place ?placeLabel ?location ?instanceLabel ?placeDescription
+            WHERE
+            {
+              wd:%s wdt:P625 ?loc .
+              SERVICE wikibase:around {
+                  ?place wdt:P625 ?location .
+                  bd:serviceParam wikibase:center ?loc .
+                  bd:serviceParam wikibase:radius "5" .
+              }
+              OPTIONAL {    ?place wdt:P31 ?instance  }
+              SERVICE wikibase:label { bd:serviceParam wikibase:language "en" }
+              BIND(geof:distance(?loc, ?location) as ?dist)
+            } ORDER BY ?dist
+            LIMIT 200
+            """%(q_id)
+            # initially 2km
+
+        #pdb.set_trace()
+        
+        response = requests.get(self.baseuri, params = {'format':'json', 'query':query})
+
+
+        ret_json = self.response_handler(response, query)
+        
+            
+
+        return ret_json
+
+
+        
+
+    def linkedgeodata_query (self, name_str):
+
+        query = """
+
+        Prefix lgdo: <http://linkedgeodata.org/ontology/>
+        Prefix geom: <http://geovocab.org/geometry#>
+        Prefix ogc: <http://www.opengis.net/ont/geosparql#>
+        Prefix owl: <http://www.w3.org/2002/07/owl#>
+        Prefix wgs84_pos: <http://www.w3.org/2003/01/geo/wgs84_pos#>
+        Prefix owl: <http://www.w3.org/2002/07/owl#>
+        Prefix gn: <http://www.geonames.org/ontology#>
+
+        Select ?s, ?lat, ?long {
+            {?s rdfs:label \"%s\";
+                wgs84_pos:lat ?lat ;
+                wgs84_pos:long ?long;
+                } 
+        } 
+        """%(name_str)
+
+        
+
+        response = requests.get(self.baseuri, params = {'format':'json', 'query':query})
+
+
+        ret_json = self.response_handler(response, query)
+
+
+        return ret_json
+
+
+    
+if __name__ == '__main__':
+    request_wrapper_wikidata = RequestWrapper(baseuri = 'https://query.wikidata.org/sparql')
+    #print(request_wrapper_wikidata.wikidata_nearby_query('Q370771'))
+    #print(request_wrapper_wikidata.wikidata_query_withinstate('San Bernardino'))
+
+    # not working now
+    print(request_wrapper_wikidata.linkedgeodata_query('San Bernardino'))

models/spabert/experiments/entity_matching/data_processing/run_linking_query.py

@@ -0,0 +1,143 @@
+#from query_wrapper import QueryWrapper
+from request_wrapper import RequestWrapper
+from get_namelist import *
+import glob
+import os
+import json
+import time
+
+DATASET_OPTIONS = ['OSM', 'OS', 'USGS', 'GB1900']
+KB_OPTIONS = ['wikidata', 'linkedgeodata']
+
+DATASET = 'USGS'
+KB = 'wikidata'
+OVERWRITE = True
+WITHIN_CA = True
+
+assert DATASET in DATASET_OPTIONS
+assert KB in KB_OPTIONS
+
+
+def process_one_namelist(sparql_wrapper, namelist, out_path):
+
+    if OVERWRITE:
+        # flush the file if it's been written
+        with open(out_path, 'w') as f:
+            f.write('')
+
+
+    for name in namelist:
+        name = name.replace('"', '')
+        name = name.strip("'")
+        if len(name) == 0:
+          continue
+        print(name)
+        mydict = dict()
+
+        if KB == 'wikidata':
+            if WITHIN_CA:
+                mydict[name] =  sparql_wrapper.wikidata_query_withinstate(name)
+            else:
+                mydict[name] =  sparql_wrapper.wikidata_query(name)
+
+        elif KB == 'linkedgeodata':
+            mydict[name] =  sparql_wrapper.linkedgeodata_query(name)
+        else:
+            raise NotImplementedError
+
+        line = json.dumps(mydict)
+        
+        with open(out_path, 'a') as f:
+            f.write(line)
+            f.write('\n')
+        time.sleep(1)
+
+
+def process_namelist_dict(sparql_wrapper, namelist_dict, out_dir):
+    i = 0
+    for map_name, namelist in namelist_dict.items():
+        # if i <=5:
+        #     i += 1
+        #     continue 
+
+        print('processing %s' %map_name)
+
+        if WITHIN_CA:
+            out_path = os.path.join(out_dir, KB + '_' + map_name + '.json')
+        else:
+            out_path = os.path.join(out_dir, KB + '_ca_' + map_name + '.json')
+
+        process_one_namelist(sparql_wrapper, namelist, out_path)
+        i+=1
+
+
+if KB == 'linkedgeodata':
+    sparql_wrapper = RequestWrapper(baseuri = 'http://linkedgeodata.org/sparql')
+elif KB == 'wikidata':
+    sparql_wrapper = RequestWrapper(baseuri = 'https://query.wikidata.org/sparql')
+else:
+    raise NotImplementedError
+
+
+
+if DATASET == 'OSM':
+    osm_dir = '../surface_form/data_sample_london/data_osm/'
+    osm_paths = glob.glob(os.path.join(osm_dir, 'embedding*.json'))
+    
+    out_path = 'outputs/'+KB+'_linking.json'
+    namelist = get_name_list_osm(osm_paths)
+    
+    print('# files',len(file_paths))
+
+    process_one_namelist(sparql_wrapper, namelist, out_path)
+    
+    
+elif DATASET == 'OS':  
+    histmap_dir = 'data/labGISReport-master/output/'
+    file_paths = glob.glob(os.path.join(histmap_dir, '10*.json'))
+    
+    out_path = 'outputs/'+KB+'_os_linking_descript.json'
+    namelist = get_name_list_usgs_od(file_paths)
+    
+    print('# files',len(file_paths))
+    
+
+    process_one_namelist(sparql_wrapper, namelist, out_path)
+    
+elif DATASET == 'USGS':
+    histmap_dir = 'data/labGISReport-master/output/'
+    file_paths = glob.glob(os.path.join(histmap_dir, 'USGS*.json'))
+    
+    if WITHIN_CA:
+        out_dir = 'outputs/' + KB +'_ca'
+    else:
+        out_dir = 'outputs/' + KB 
+    namelist_dict = get_name_list_usgs_od_per_map(file_paths)
+
+    if not os.path.isdir(out_dir):
+        os.makedirs(out_dir)
+    
+    print('# files',len(file_paths))
+
+    process_namelist_dict(sparql_wrapper, namelist_dict, out_dir)
+
+elif DATASET == 'GB1900':
+
+    file_path = 'data/GB1900_gazetteer_abridged_july_2018/gb1900_abridged.csv'
+    out_path = 'outputs/'+KB+'_gb1900_linking_descript.json'
+    namelist = get_name_list_gb1900(file_path)
+
+
+    process_one_namelist(sparql_wrapper, namelist, out_path)
+    
+else:
+    raise NotImplementedError
+
+
+
+
+#namelist = namelist[730:]   #for GB1900
+
+
+
+print('done')

models/spabert/experiments/entity_matching/data_processing/run_map_neighbor_query.py

@@ -0,0 +1,123 @@
+from query_wrapper import QueryWrapper
+from request_wrapper import RequestWrapper
+import glob
+import os
+import json
+import time
+import random
+
+DATASET_OPTIONS = ['OSM', 'OS', 'USGS', 'GB1900']
+KB_OPTIONS = ['wikidata', 'linkedgeodata']
+
+dataset = 'USGS'
+kb = 'wikidata'
+overwrite = False
+
+assert dataset in DATASET_OPTIONS
+assert kb in KB_OPTIONS
+
+if dataset == 'OSM':
+    raise NotImplementedError
+    
+elif dataset == 'OS':  
+    raise NotImplementedError
+    
+elif dataset == 'USGS':
+
+    candidate_file_paths = glob.glob('outputs/alignment_dir/wikidata_USGS*.json')
+    candidate_file_paths = sorted(candidate_file_paths)
+    
+    out_dir = 'outputs/wikidata_neighbors/'
+
+    if not os.path.isdir(out_dir):
+        os.makedirs(out_dir)
+
+elif dataset == 'GB1900':
+
+    raise NotImplementedError
+    
+else:
+    raise NotImplementedError
+
+
+if kb == 'linkedgeodata':
+    sparql_wrapper = QueryWrapper(baseuri = 'http://linkedgeodata.org/sparql')
+elif kb == 'wikidata':
+    #sparql_wrapper = QueryWrapper(baseuri = 'https://query.wikidata.org/sparql')
+    sparql_wrapper = RequestWrapper(baseuri = 'https://query.wikidata.org/sparql')
+else:
+    raise NotImplementedError
+
+start_map = 6   # 6 
+start_line = 4    # 4 
+
+for candiate_file_path in candidate_file_paths[start_map:]:
+    map_name = os.path.basename(candiate_file_path).split('wikidata_')[1]
+    out_path = os.path.join(out_dir, 'wikidata_' + map_name)
+
+    with open(candiate_file_path, 'r') as f:
+        cand_data = f.readlines()
+
+    with open(out_path, 'a') as out_f:
+        for line in cand_data[start_line:]:
+            line_dict = json.loads(line)
+            ret_line_dict = dict()
+            for key, value in line_dict.items(): # actually just one pair
+
+                print(key)
+
+                place_name = key 
+                for cand_entity in value:
+                    time.sleep(2)
+                    q_id = cand_entity['item']['value'].split('/')[-1]
+                    response = sparql_wrapper.wikidata_nearby_query(str(q_id))
+
+                    if place_name in ret_line_dict:
+                        ret_line_dict[place_name].append(response)
+                    else:
+                        ret_line_dict[place_name] = [response]
+
+                    #time.sleep(random.random()*6)
+                    
+
+
+            out_f.write(json.dumps(ret_line_dict))
+            out_f.write('\n')
+
+    print('finished with ',candiate_file_path)
+    break
+                    
+print('done')
+
+'''
+        {"Martin": [{"item": {"type": "uri", "value": "http://www.wikidata.org/entity/Q27001"}, "coordinates": {"datatype": "http://www.opengis.net/ont/geosparql#wktLiteral", "type": "literal", "value": "Point(18.921388888 49.063611111)"}, "itemDescription": {"xml:lang": "en", "type": "literal", "value": "city in northern Slovakia"}}, {"item": {"type": "uri", "value": "http://www.wikidata.org/entity/Q281028"}, "coordinates": {"datatype": "http://www.opengis.net/ont/geosparql#wktLiteral", "type": "literal", "value": "Point(-101.734166666 43.175)"}, "itemDescription": {"xml:lang": "en", "type": "literal", "value": "city in and county seat of Bennett County, South Dakota, United States"}}, {"item": {"type": "uri", "value": "http://www.wikidata.org/entity/Q761390"}, "coordinates": {"datatype": "http://www.opengis.net/ont/geosparql#wktLiteral", "type": "literal", "value": "Point(1.3394 51.179)"}, "itemDescription": {"xml:lang": "en", "type": "literal", "value": "hamlet in Kent"}}, {"item": {"type": "uri", "value": "http://www.wikidata.org/entity/Q2177502"}, "coordinates": {"datatype": "http://www.opengis.net/ont/geosparql#wktLiteral", "type": "literal", "value": "Point(-100.115 47.826666666)"}, "itemDescription": {"xml:lang": "en", "type": "literal", "value": "city in North Dakota"}}, {"item": {"type": "uri", "value": "http://www.wikidata.org/entity/Q2454021"}, "coordinates": {"datatype": "http://www.opengis.net/ont/geosparql#wktLiteral", "type": "literal", "value": "Point(-85.64168 42.53698)"}, "itemDescription": {"xml:lang": "en", "type": "literal", "value": "village in Michigan, USA"}}, {"item": {"type": "uri", "value": "http://www.wikidata.org/entity/Q2481111"}, "coordinates": {"datatype": "http://www.opengis.net/ont/geosparql#wktLiteral", "type": "literal", "value": "Point(-93.21833 32.09917)"}, "itemDescription": {"xml:lang": "en", "type": "literal", "value": "village in Red River Parish, Louisiana, United States"}}, {"item": {"type": "uri", "value": "http://www.wikidata.org/entity/Q2635473"}, "coordinates": {"datatype": "http://www.opengis.net/ont/geosparql#wktLiteral", "type": "literal", "value": "Point(-82.75944 37.56778)"}, "itemDescription": {"xml:lang": "en", "type": "literal", "value": "city in Kentucky, USA"}}, {"item": {"type": "uri", "value": "http://www.wikidata.org/entity/Q2679547"}, "coordinates": {"datatype": "http://www.opengis.net/ont/geosparql#wktLiteral", "type": "literal", "value": "Point(-1.9041 50.9759)"}, "itemDescription": {"xml:lang": "en", "type": "literal", "value": "village in Hampshire, England, United Kingdom"}}, {"item": {"type": "uri", "value": "http://www.wikidata.org/entity/Q2780056"}, "coordinates": {"datatype": "http://www.opengis.net/ont/geosparql#wktLiteral", "type": "literal", "value": "Point(-88.851666666 36.341944444)"}, "itemDescription": {"xml:lang": "en", "type": "literal", "value": "city in Tennessee, USA"}}, {"item": {"type": "uri", "value": "http://www.wikidata.org/entity/Q3261150"}, "coordinates": {"datatype": "http://www.opengis.net/ont/geosparql#wktLiteral", "type": "literal", "value": "Point(-83.18556 34.48639)"}, "itemDescription": {"xml:lang": "en", "type": "literal", "value": "town in Franklin and Stephens Counties, Georgia, USA"}}, {"item": {"type": "uri", "value": "http://www.wikidata.org/entity/Q6002227"}, "coordinates": {"datatype": "http://www.opengis.net/ont/geosparql#wktLiteral", "type": "literal", "value": "Point(-101.709 41.2581)"}, "itemDescription": {"xml:lang": "en", "type": "literal", "value": "census-designated place in Nebraska, United States"}}, {"item": {"type": "uri", "value": "http://www.wikidata.org/entity/Q6774807"}, "coordinates": {"datatype": "http://www.opengis.net/ont/geosparql#wktLiteral", "type": "literal", "value": "Point(-82.1906 29.2936)"}, "itemDescription": {"xml:lang": "en", "type": "literal", "value": "unincorporated community in Marion County, Florida, United States"}}, {"item": {"type": "uri", "value": "http://www.wikidata.org/entity/Q6774809"}, "coordinates": {"datatype": "http://www.opengis.net/ont/geosparql#wktLiteral", "type": "literal", "value": "Point(-83.336666666 41.5575)"}, "itemDescription": {"xml:lang": "en", "type": "literal", "value": "unincorporated community in Ohio, United States"}}, {"item": {"type": "uri", "value": "http://www.wikidata.org/entity/Q6774810"}, "coordinates": {"datatype": "http://www.opengis.net/ont/geosparql#wktLiteral", "type": "literal", "value": "Point(116.037 -32.071)"}, "itemDescription": {"xml:lang": "en", "type": "literal", "value": "suburb of Perth, Western Australia"}}, {"item": {"type": "uri", "value": "http://www.wikidata.org/entity/Q9029707"}, "coordinates": {"datatype": "http://www.opengis.net/ont/geosparql#wktLiteral", "type": "literal", "value": "Point(-81.476388888 33.069166666)"}, "itemDescription": {"xml:lang": "en", "type": "literal", "value": "unincorporated community in South Carolina, United States"}}, {"item": {"type": "uri", "value": "http://www.wikidata.org/entity/Q11770660"}, "coordinates": {"datatype": "http://www.opengis.net/ont/geosparql#wktLiteral", "type": "literal", "value": "Point(-0.325391 53.1245)"}, "itemDescription": {"xml:lang": "en", "type": "literal", "value": "village and civil parish in Lincolnshire, UK"}}, {"item": {"type": "uri", "value": "http://www.wikidata.org/entity/Q14692833"}, "coordinates": {"datatype": "http://www.opengis.net/ont/geosparql#wktLiteral", "type": "literal", "value": "Point(-93.5444 47.4894)"}, "itemDescription": {"xml:lang": "en", "type": "literal", "value": "unincorporated community in Itasca County, Minnesota"}}, {"item": {"type": "uri", "value": "http://www.wikidata.org/entity/Q14714180"}, "coordinates": {"datatype": "http://www.opengis.net/ont/geosparql#wktLiteral", "type": "literal", "value": "Point(-79.0889 39.2242)"}, "itemDescription": {"xml:lang": "en", "type": "literal", "value": "unincorporated community in Grant County, West Virginia"}}, {"item": {"type": "uri", "value": "http://www.wikidata.org/entity/Q18496647"}, "coordinates": {"datatype": "http://www.opengis.net/ont/geosparql#wktLiteral", "type": "literal", "value": "Point(11.95941 46.34585)"}, "itemDescription": {"xml:lang": "en", "type": "literal", "value": "human settlement in Italy"}}, {"item": {"type": "uri", "value": "http://www.wikidata.org/entity/Q20949553"}, "coordinates": {"datatype": "http://www.opengis.net/ont/geosparql#wktLiteral", "type": "literal", "value": "Point(22.851944444 41.954444444)"}, "itemDescription": {"xml:lang": "en", "type": "literal", "value": "mountain in Republic of Macedonia"}}, {"item": {"type": "uri", "value": "http://www.wikidata.org/entity/Q24065096"}, "coordinates": {"datatype": "http://www.opengis.net/ont/geosparql#wktLiteral", "type": "literal", "value": "<http://www.wikidata.org/entity/Q111> Point(290.75 -21.34)"}, "itemDescription": {"xml:lang": "en", "type": "literal", "value": "crater on Mars"}}, {"item": {"type": "uri", "value": "http://www.wikidata.org/entity/Q26300074"}, "coordinates": {"datatype": "http://www.opengis.net/ont/geosparql#wktLiteral", "type": "literal", "value": "Point(-0.97879914 51.74729077)"}, "itemDescription": {"xml:lang": "en", "type": "literal", "value": "Thame, South Oxfordshire, Oxfordshire, OX9"}}, {"item": {"type": "uri", "value": "http://www.wikidata.org/entity/Q27988822"}, "coordinates": {"datatype": "http://www.opengis.net/ont/geosparql#wktLiteral", "type": "literal", "value": "Point(-87.67361111 38.12444444)"}, "itemDescription": {"xml:lang": "en", "type": "literal", "value": "human settlement in Vanderburgh County, Indiana, United States"}}, {"item": {"type": "uri", "value": "http://www.wikidata.org/entity/Q27995389"}, "coordinates": {"datatype": "http://www.opengis.net/ont/geosparql#wktLiteral", "type": "literal", "value": "Point(-121.317 47.28)"}, "itemDescription": {"xml:lang": "en", "type": "literal", "value": "human settlement in Washington, United States of America"}}, {"item": {"type": "uri", "value": "http://www.wikidata.org/entity/Q28345614"}, "coordinates": {"datatype": "http://www.opengis.net/ont/geosparql#wktLiteral", "type": "literal", "value": "Point(-76.8 48.5)"}, "itemDescription": {"xml:lang": "en", "type": "literal", "value": "human settlement in Senneterre, Quebec, Canada"}}, {"item": {"type": "uri", "value": "http://www.wikidata.org/entity/Q30626037"}, "coordinates": {"datatype": "http://www.opengis.net/ont/geosparql#wktLiteral", "type": "literal", "value": "Point(-79.90972222 39.80638889)"}, "itemDescription": {"xml:lang": "en", "type": "literal", "value": "human settlement in United States of America"}}, {"item": {"type": "uri", "value": "http://www.wikidata.org/entity/Q61038281"}, "coordinates": {"datatype": "http://www.opengis.net/ont/geosparql#wktLiteral", "type": "literal", "value": "Point(-91.13 49.25)"}, "itemDescription": {"xml:lang": "en", "type": "literal", "value": "Meteorological Service of Canada's station for Martin (MSC ID: 6035000), Ontario, Canada"}}, {"item": {"type": "uri", "value": "http://www.wikidata.org/entity/Q63526691"}, "coordinates": {"datatype": "http://www.opengis.net/ont/geosparql#wktLiteral", "type": "literal", "value": "Point(152.7011111 -29.881666666)"}, "itemDescription": {"xml:lang": "en", "type": "literal", "value": "Parish of Fitzroy County, New South Wales, Australia"}}, {"item": {"type": "uri", "value": "http://www.wikidata.org/entity/Q63526695"}, "coordinates": {"datatype": "http://www.opengis.net/ont/geosparql#wktLiteral", "type": "literal", "value": "Point(148.1011111 -33.231666666)"}, "itemDescription": {"xml:lang": "en", "type": "literal", "value": "Parish of Ashburnham County, New South Wales, Australia"}}, {"item": {"type": "uri", "value": "http://www.wikidata.org/entity/Q96149222"}, "coordinates": {"datatype": "http://www.opengis.net/ont/geosparql#wktLiteral", "type": "literal", "value": "Point(14.725573779 48.76768332)"}}, {"item": {"type": "uri", "value": "http://www.wikidata.org/entity/Q96158116"}, "coordinates": {"datatype": "http://www.opengis.net/ont/geosparql#wktLiteral", "type": "literal", "value": "Point(15.638358708 49.930136157)"}}, {"item": {"type": "uri", "value": "http://www.wikidata.org/entity/Q103777024"}, "coordinates": {"datatype": "http://www.opengis.net/ont/geosparql#wktLiteral", "type": "literal", "value": "Point(-3.125028822 58.694748091)"}, "itemDescription": {"xml:lang": "en", "type": "literal", "value": "Shipwreck off the Scottish Coast, imported from Canmore Nov 2020"}}, {"item": {"type": "uri", "value": "http://www.wikidata.org/entity/Q107077206"}, "coordinates": {"datatype": "http://www.opengis.net/ont/geosparql#wktLiteral", "type": "literal", "value": "Point(11.688165 46.582982)"}}]}
+
+
+if overwrite:
+    # flush the file if it's been written
+    with open(out_path, 'w') as f:
+        f.write('')
+
+for name in namelist:
+    name = name.replace('"', '')
+    name = name.strip("'")
+    if len(name) == 0:
+      continue
+    print(name)
+    mydict = dict()
+    mydict[name] =  sparql_wrapper.wikidata_query(name)
+    line = json.dumps(mydict)
+    #print(line)
+    with open(out_path, 'a') as f:
+        f.write(line)
+        f.write('\n')
+    time.sleep(1)
+
+print('done')
+
+
+{"info": {"name": "10TH ST", "geometry": [4193.118085062303, -831.274950414831]}, 
+"neighbor_info": 
+{"name_list": ["BM 107", "PALM AVE", "WT", "Hidalgo Sch", "PO", "MAIN ST", "BRYANT CANAL", "BM 123", "Oakley Sch", "BRAWLEY", "Witter Sch", "BM 104", "Pistol Range", "Reid Sch", "STANLEY", "MUNICIPAL AIRPORT", "WESTERN AVE", "CANAL", "Riverview Cem", "BEST CANAL"], 
+"geometry_list": [[4180.493095652702, -836.0635465095995], [4240.450935702045, -855.345637906981], [4136.084840542623, -917.7895986922882], [4150.386997979736, -948.7258091165079], [4056.955267048625, -847.1018277439381], [4008.112642182582, -849.089249977583], [4124.177447575567, -1004.0706369942257], [4145.382175508665, -626.1608201557082], [4398.137868976953, -764.1087236140554], [4221.1546492913285, -1062.5745271963772], [4015.203890157584, -985.0178210457995], [3989.2345421184878, -948.9340389243871], [4385.585449075614, -660.4590917125413], [3936.505159635338, -803.6822663422273], [3960.1233867112846, -686.7988766730389], [4409.714306709143, -600.6633389979504], [3871.2873706574037, -832.0785684368772], [4304.899727301024, -524.472390102557], [3955.640201659347, -578.5544271698675], [4075.8524354668034, -1183.5837385075774]]}}
+'''

models/spabert/experiments/entity_matching/data_processing/run_query_sample.py

@@ -0,0 +1,22 @@
+from query_wrapper import QueryWrapper
+from get_namelist import *
+import glob
+import os
+import json
+import time
+
+
+#sparql_wrapper_linkedgeo = QueryWrapper(baseuri = 'http://linkedgeodata.org/sparql')
+
+#print(sparql_wrapper_linkedgeo.linkedgeodata_query('Los Angeles'))
+
+
+sparql_wrapper_wikidata = QueryWrapper(baseuri = 'https://query.wikidata.org/sparql')
+
+#print(sparql_wrapper_wikidata.wikidata_query('Los Angeles'))
+
+#time.sleep(3)
+
+#print(sparql_wrapper_wikidata.wikidata_nearby_query('Q370771'))
+print(sparql_wrapper_wikidata.wikidata_nearby_query('Q97625145'))
+

models/spabert/experiments/entity_matching/data_processing/run_wikidata_neighbor_query.py

@@ -0,0 +1,31 @@
+import pandas as pd
+import json
+from request_wrapper import RequestWrapper
+import time
+import pdb 
+
+start_idx = 17335
+wikidata_sample30k_path = 'wikidata_sample30k/wikidata_30k.json'
+out_path = 'wikidata_sample30k/wikidata_30k_neighbor.json'
+
+#with open(out_path, 'w') as out_f:
+#    pass 
+
+sparql_wrapper = RequestWrapper(baseuri = 'https://query.wikidata.org/sparql')
+
+df= pd.read_json(wikidata_sample30k_path)
+df = df[start_idx:]
+
+print('length of df:', len(df))
+
+for index, record in df.iterrows():
+    print(index)
+    uri = record.results['item']['value']
+    q_id = uri.split('/')[-1]
+    response = sparql_wrapper.wikidata_nearby_query(str(q_id))
+    time.sleep(1)
+    with open(out_path, 'a') as out_f:
+        out_f.write(json.dumps(response))
+        out_f.write('\n')
+    
+

models/spabert/experiments/entity_matching/data_processing/samples.sparql

@@ -0,0 +1,22 @@
+'''
+Entities near xxx within 1km
+SELECT ?place ?placeLabel ?location ?instanceLabel ?placeDescription
+WHERE
+{
+  wd:Q9188 wdt:P625 ?loc .
+  SERVICE wikibase:around {
+      ?place wdt:P625 ?location .
+      bd:serviceParam wikibase:center ?loc .
+      bd:serviceParam wikibase:radius "1" .
+  }
+  OPTIONAL {    ?place wdt:P31 ?instance  }
+  SERVICE wikibase:label { bd:serviceParam wikibase:language "en" }
+  BIND(geof:distance(?loc, ?location) as ?dist)
+} ORDER BY ?dist
+'''
+
+
+SELECT distinct ?item ?itemLabel WHERE {
+  ?item wdt:P625 ?geo . 
+SERVICE wikibase:label { bd:serviceParam wikibase:language "[AUTO_LANGUAGE],en". }
+}

models/spabert/experiments/entity_matching/data_processing/select_ambi.py

@@ -0,0 +1,18 @@
+import json
+
+json_file = 'entity_linking/outputs/wikidata_usgs_linking_descript.json'
+
+with open(json_file, 'r') as f:
+	data = f.readlines()
+
+num_ambi = 0
+for line in data:
+	line_dict = json.loads(line)
+	for key,value in line_dict.items():
+		len_value = len(value)
+		if len_value < 2:
+			continue
+		else:
+			num_ambi += 1
+			print(key)
+print(num_ambi)

models/spabert/experiments/entity_matching/src/evaluation-mrr.py

@@ -0,0 +1,260 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+
+import sys
+import os
+import glob
+import json
+import numpy as np
+import pandas as pd
+import pdb
+
+
+prediction_dir = sys.argv[1]
+
+print(prediction_dir)
+
+gt_dir = '../data_processing/outputs/alignment_gt_dir/'
+prediction_path_list = sorted(os.listdir(prediction_dir))
+
+DISPLAY = False
+DETAIL = False
+
+if DISPLAY:
+    from IPython.display import display
+
+def recall_at_k_all_map(all_rank_list, k = 1):
+    
+    rank_list = [item for sublist in all_rank_list for item in sublist]
+    total_query = len(rank_list)
+    prec = np.sum(np.array(rank_list)<=k)
+    prec = 1.0 * prec / total_query
+    
+    return prec
+
+def recall_at_k_permap(all_rank_list, k = 1):
+    
+    prec_list = []
+    for rank_list in all_rank_list:
+        total_query = len(rank_list)
+        prec = np.sum(np.array(rank_list)<=k)
+        prec = 1.0 * prec / total_query
+        prec_list.append(prec)
+    
+    return prec_list
+
+
+
+
+def reciprocal_rank(all_rank_list):
+    
+    recip_list = [1./rank for rank in all_rank_list]
+    mean_recip = np.mean(recip_list)
+    
+    return mean_recip, recip_list
+
+
+
+
+count_hist_list = []
+
+all_rank_list = []
+
+all_recip_list = []
+
+permap_recip_list = []
+
+for map_path in prediction_path_list:
+    
+    pred_path = os.path.join(prediction_dir, map_path)
+    gt_path = os.path.join(gt_dir, map_path.split('.json')[0] + '.csv')
+    
+    if DETAIL:
+        print(pred_path)
+    
+    
+    with open(gt_path, 'r') as f:
+        gt_data = f.readlines()
+        
+    gt_dict = dict()
+    for line in gt_data:
+        line = line.split(',')
+        pivot_name = line[0]
+        gt_uri = line[1]
+        gt_dict[pivot_name] = gt_uri
+        
+    rank_list = []
+    pivot_name_list = []
+    with open(pred_path, 'r') as f:
+        pred_data = f.readlines()
+        for line in pred_data:
+            pred_dict = json.loads(line)
+            #print(pred_dict.keys())
+            pivot_name = pred_dict['pivot_name']
+            sorted_match_uri = pred_dict['sorted_match_uri']
+            #sorted_match_des = pred_dict['sorted_match_des']
+            sorted_sim_matrix = pred_dict['sorted_sim_matrix']
+            
+            
+            total = len(sorted_match_uri)
+            if total == 1: 
+                continue
+            
+            if pivot_name in gt_dict:
+                
+                gt_uri = gt_dict[pivot_name]
+                
+                try:
+                    assert gt_uri in sorted_match_uri
+                except Exception as e:
+                    #print(e)
+                    continue
+                
+                pivot_name_list.append(pivot_name)
+                count_hist_list.append(total)
+                rank = sorted_match_uri.index(gt_uri) +1
+                
+                rank_list.append(rank)
+                #print(rank,'/',total)
+                
+    all_rank_list.append(rank_list)
+    
+    mean_recip, recip_list = reciprocal_rank(rank_list)
+
+    all_recip_list.extend(recip_list)
+    permap_recip_list.append(recip_list)
+    
+    d = {'pivot': pivot_name_list + ['AVG'], 'rank':rank_list + [' '] ,'recip rank': recip_list + [str(mean_recip)]}
+    if DETAIL:
+        print(pivot_name_list, rank_list, recip_list)
+    
+    if DISPLAY:
+        df = pd.DataFrame(data=d)
+
+        display(df)
+    
+    
+
+print('all mrr, micro', np.mean(all_recip_list))
+
+
+if DETAIL:
+
+    len(rank_list)
+
+
+
+    print(recall_at_k_all_map(all_rank_list, k = 1))
+    print(recall_at_k_all_map(all_rank_list, k = 2))
+    print(recall_at_k_all_map(all_rank_list, k = 5))
+    print(recall_at_k_all_map(all_rank_list, k = 10))
+
+
+    print(prediction_path_list)
+
+
+prec_list_1 = recall_at_k_permap(all_rank_list, k = 1)
+prec_list_2 = recall_at_k_permap(all_rank_list, k = 2)
+prec_list_5 = recall_at_k_permap(all_rank_list, k = 5)
+prec_list_10 = recall_at_k_permap(all_rank_list, k = 10)
+
+if DETAIL:
+
+    print(np.mean(prec_list_1))
+    print(prec_list_1)
+    print('\n')
+
+    print(np.mean(prec_list_2))
+    print(prec_list_2)
+    print('\n')
+
+    print(np.mean(prec_list_5))
+    print(prec_list_5)
+    print('\n')
+
+    print(np.mean(prec_list_10))
+    print(prec_list_10)
+    print('\n')
+
+
+
+
+
+
+
+
+import pandas as pd
+
+
+
+map_name_list = [name.split('.json')[0].split('USGS-')[1] for name in prediction_path_list]
+d = {'map_name': map_name_list,'recall@1': prec_list_1, 'recall@2': prec_list_2, 'recall@5': prec_list_5, 'recall@10': prec_list_10 }
+df = pd.DataFrame(data=d)
+
+
+if DETAIL:
+    print(df)
+
+
+
+
+
+category = ['15-CA','30-CA','60-CA']
+col_1 = [np.mean(prec_list_1[0:4]), np.mean(prec_list_1[4:9]), np.mean(prec_list_1[9:])]
+col_2 = [np.mean(prec_list_2[0:4]), np.mean(prec_list_2[4:9]), np.mean(prec_list_2[9:])]
+col_3 = [np.mean(prec_list_5[0:4]), np.mean(prec_list_5[4:9]), np.mean(prec_list_5[9:])]
+col_4 = [np.mean(prec_list_10[0:4]), np.mean(prec_list_10[4:9]), np.mean(prec_list_10[9:])]
+
+
+
+mrr_15 = permap_recip_list[0] + permap_recip_list[1] + permap_recip_list[2] + permap_recip_list[3]
+mrr_30 = permap_recip_list[4] + permap_recip_list[5] + permap_recip_list[6] + permap_recip_list[7] + permap_recip_list[8] 
+mrr_60 = permap_recip_list[9] + permap_recip_list[10] + permap_recip_list[11] + permap_recip_list[12] + permap_recip_list[13] 
+
+
+
+column_5 = [np.mean(mrr_15), np.mean(mrr_30), np.mean(mrr_60)]
+
+
+d = {'map set': category, 'mrr': column_5, 'prec@1': col_1, 'prec@2': col_2, 'prec@5': col_3, 'prec@10': col_4 }
+df = pd.DataFrame(data=d)
+
+print(df)
+
+
+
+
+print('all mrr, micro', np.mean(all_recip_list))
+
+print('\n')
+
+
+
+print(recall_at_k_all_map(all_rank_list, k = 1))
+print(recall_at_k_all_map(all_rank_list, k = 2))
+print(recall_at_k_all_map(all_rank_list, k = 5))
+print(recall_at_k_all_map(all_rank_list, k = 10))
+
+
+
+
+if DISPLAY:
+
+    import seaborn 
+
+    p = seaborn.histplot(data = count_hist_list, color = 'blue', alpha=0.2)
+    p.set_xlabel("Number of Candiates")
+    p.set_title("Candidate Distribution in USGS")
+
+
+
+
+    len(count_hist_list)
+
+
+
+
+
+
+

models/spabert/experiments/entity_matching/src/linking_ablation.py

@@ -0,0 +1,228 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+import sys
+import os
+import numpy as np
+import pdb
+import json
+import scipy.spatial as sp
+import argparse
+
+
+import torch
+from torch.utils.data import DataLoader
+
+from transformers import AdamW
+from transformers import BertTokenizer
+from tqdm import tqdm  # for our progress bar
+
+sys.path.append('../../../')
+from datasets.usgs_os_sample_loader import USGS_MapDataset
+from datasets.wikidata_sample_loader import Wikidata_Geocoord_Dataset, Wikidata_Random_Dataset
+from models.spatial_bert_model import SpatialBertModel
+from models.spatial_bert_model import SpatialBertConfig
+from utils.find_closest import find_ref_closest_match, sort_ref_closest_match
+from utils.common_utils import load_spatial_bert_pretrained_weights, get_spatialbert_embedding, get_bert_embedding, write_to_csv
+from utils.baseline_utils import get_baseline_model
+
+from transformers import BertModel
+
+sys.path.append('/home/zekun/spatial_bert/spatial_bert/datasets')
+from dataset_loader import SpatialDataset
+from osm_sample_loader import PbfMapDataset
+
+
+
+MODEL_OPTIONS = ['spatial_bert-base','spatial_bert-large', 'bert-base','bert-large','roberta-base','roberta-large',
+'spanbert-base','spanbert-large','luke-base','luke-large',
+'simcse-bert-base','simcse-bert-large','simcse-roberta-base','simcse-roberta-large']
+
+
+CANDSET_MODES = ['all_map'] # candidate set is constructed based on all maps or one map
+
+def recall_at_k(rank_list, k = 1):
+    
+    total_query = len(rank_list)
+    recall = np.sum(np.array(rank_list)<=k)
+    recall = 1.0 * recall / total_query
+    
+    return recall
+
+def reciprocal_rank(all_rank_list):
+    
+    recip_list = [1./rank for rank in all_rank_list]
+    mean_recip = np.mean(recip_list)
+    
+    return mean_recip, recip_list
+
+def link_to_itself(source_embedding_ogc_list, target_embedding_ogc_list):
+
+    source_emb_list = [source_dict['emb'] for source_dict in source_embedding_ogc_list]
+    source_ogc_list = [source_dict['ogc_fid'] for source_dict in source_embedding_ogc_list]
+    
+    target_emb_list = [target_dict['emb'] for target_dict in target_embedding_ogc_list]
+    target_ogc_list = [target_dict['ogc_fid'] for target_dict in target_embedding_ogc_list]
+
+    rank_list = []
+    for source_emb, source_ogc in zip(source_emb_list, source_ogc_list):
+        sim_matrix = 1 - sp.distance.cdist(np.array(target_emb_list), np.array([source_emb]), 'cosine')
+        closest_match_ogc = sort_ref_closest_match(sim_matrix, target_ogc_list)
+
+        closest_match_ogc = [a[0] for a in closest_match_ogc]
+        rank = closest_match_ogc.index(source_ogc) +1
+        rank_list.append(rank)
+
+    
+    mean_recip, recip_list = reciprocal_rank(rank_list)
+    r1 = recall_at_k(rank_list, k = 1)
+    r5 = recall_at_k(rank_list, k = 5)
+    r10 = recall_at_k(rank_list, k = 10)
+    
+    return mean_recip , r1, r5, r10
+
+def get_embedding_and_ogc(dataset, model_name, model):
+    dict_list = []
+
+    for source in dataset:
+        if model_name == 'spatial_bert-base' or model_name ==  'spatial_bert-large':
+            source_emb = get_spatialbert_embedding(source, model)
+        else:
+            source_emb = get_bert_embedding(source, model)
+            
+        source_dict = {}
+        source_dict['emb'] = source_emb
+        source_dict['ogc_fid'] = source['ogc_fid']
+        #wikidata_dict['wikidata_des_list'] = [wikidata_cand['description']]
+
+        dict_list.append(source_dict)
+
+    return dict_list
+
+
+def entity_linking_func(args):
+
+    model_name = args.model_name
+    candset_mode = args.candset_mode
+
+    distance_norm_factor = args.distance_norm_factor
+    spatial_dist_fill= args.spatial_dist_fill
+    sep_between_neighbors = args.sep_between_neighbors
+
+    spatial_bert_weight_dir = args.spatial_bert_weight_dir
+    spatial_bert_weight_name = args.spatial_bert_weight_name
+
+    if_no_spatial_distance = args.no_spatial_distance
+    random_remove_neighbor = args.random_remove_neighbor
+
+
+    assert model_name in MODEL_OPTIONS
+    assert candset_mode in CANDSET_MODES
+
+
+    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+
+    
+        
+    if model_name == 'spatial_bert-base':
+        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+        
+        config = SpatialBertConfig()
+        model = SpatialBertModel(config)
+
+        model.to(device)
+        model.eval()
+        
+        # load pretrained weights
+        weight_path = os.path.join(spatial_bert_weight_dir, spatial_bert_weight_name)
+        model = load_spatial_bert_pretrained_weights(model, weight_path)
+
+    elif model_name == 'spatial_bert-large':
+        tokenizer = BertTokenizer.from_pretrained('bert-large-uncased')
+        
+        config = SpatialBertConfig(hidden_size = 1024, intermediate_size = 4096, num_attention_heads=16, num_hidden_layers=24)
+        model = SpatialBertModel(config)
+
+        model.to(device)
+        model.eval()
+        
+        # load pretrained weights
+        weight_path = os.path.join(spatial_bert_weight_dir, spatial_bert_weight_name)
+        model = load_spatial_bert_pretrained_weights(model, weight_path)
+
+    else:
+        model, tokenizer = get_baseline_model(model_name)
+        model.to(device)
+        model.eval()
+
+    source_file_path = '../data/osm-point-minnesota-full.json'
+    source_dataset = PbfMapDataset(data_file_path = source_file_path, 
+                                        tokenizer = tokenizer, 
+                                        max_token_len = 512, 
+                                        distance_norm_factor = distance_norm_factor, 
+                                        spatial_dist_fill = spatial_dist_fill, 
+                                        with_type = False,
+                                        sep_between_neighbors = sep_between_neighbors, 
+                                        mode = None,
+                                        random_remove_neighbor = random_remove_neighbor,
+                                        )
+
+    target_dataset = PbfMapDataset(data_file_path = source_file_path, 
+                                        tokenizer = tokenizer, 
+                                        max_token_len = 512, 
+                                        distance_norm_factor = distance_norm_factor, 
+                                        spatial_dist_fill = spatial_dist_fill, 
+                                        with_type = False,
+                                        sep_between_neighbors = sep_between_neighbors, 
+                                        mode = None,
+                                        random_remove_neighbor = 0., # keep all
+                                        )
+
+    # process candidates for each phrase
+    
+
+    source_embedding_ogc_list = get_embedding_and_ogc(source_dataset, model_name, model)
+    target_embedding_ogc_list = get_embedding_and_ogc(target_dataset, model_name, model)
+
+        
+    mean_recip , r1, r5, r10 = link_to_itself(source_embedding_ogc_list, target_embedding_ogc_list)
+    print('\n')
+    print(random_remove_neighbor, mean_recip , r1, r5, r10)
+
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--model_name', type=str, default='spatial_bert-base')
+    parser.add_argument('--candset_mode', type=str, default='all_map')
+
+    parser.add_argument('--distance_norm_factor', type=float, default = 0.0001)
+    parser.add_argument('--spatial_dist_fill', type=float, default = 20)
+                       
+    parser.add_argument('--sep_between_neighbors', default=False, action='store_true')
+    parser.add_argument('--no_spatial_distance',  default=False, action='store_true')
+
+    parser.add_argument('--spatial_bert_weight_dir', type = str, default = None)
+    parser.add_argument('--spatial_bert_weight_name', type = str, default = None)
+
+    parser.add_argument('--random_remove_neighbor', type = float, default = 0.)
+    
+                        
+    args = parser.parse_args()
+    # print('\n')
+    # print(args)
+    # print('\n')
+
+    entity_linking_func(args)
+
+    # CUDA_VISIBLE_DEVICES='1' python3 linking_ablation.py  --sep_between_neighbors --model_name='spatial_bert-base'  --spatial_bert_weight_dir='/data/zekun/spatial_bert_weights/typing_lr5e-05_sep_bert-base_nofreeze_london_california_bsize12/ep0_iter06000_0.2936/' --spatial_bert_weight_name='keeppos_ep0_iter02000_0.4879.pth' --random_remove_neighbor=0.1
+
+
+    # CUDA_VISIBLE_DEVICES='1' python3 linking_ablation.py  --sep_between_neighbors --model_name='spatial_bert-large'  --spatial_bert_weight_dir='/data/zekun/spatial_bert_weights/typing_lr1e-06_sep_bert-large_nofreeze_london_california_bsize12/ep2_iter02000_0.3921/' --spatial_bert_weight_name='keeppos_ep8_iter03568_0.2661_val0.2284.pth' --random_remove_neighbor=0.1
+
+
+if __name__ == '__main__':
+
+    main()
+
+    

models/spabert/experiments/entity_matching/src/unsupervised_wiki_location_allcand.py

@@ -0,0 +1,329 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+import sys
+import os
+import numpy as np
+import pdb
+import json
+import scipy.spatial as sp
+import argparse
+
+
+import torch
+from torch.utils.data import DataLoader
+#from transformers.models.bert.modeling_bert import BertForMaskedLM
+
+from transformers import AdamW
+from transformers import BertTokenizer
+from tqdm import tqdm  # for our progress bar
+
+sys.path.append('../../../')
+from datasets.usgs_os_sample_loader import USGS_MapDataset
+from datasets.wikidata_sample_loader import Wikidata_Geocoord_Dataset, Wikidata_Random_Dataset
+from models.spatial_bert_model import SpatialBertModel
+from models.spatial_bert_model import SpatialBertConfig
+#from models.spatial_bert_model import  SpatialBertForMaskedLM
+from utils.find_closest import find_ref_closest_match, sort_ref_closest_match
+from utils.common_utils import load_spatial_bert_pretrained_weights, get_spatialbert_embedding, get_bert_embedding, write_to_csv
+from utils.baseline_utils import get_baseline_model
+
+from transformers import BertModel
+
+
+MODEL_OPTIONS = ['spatial_bert-base','spatial_bert-large', 'bert-base','bert-large','roberta-base','roberta-large',
+'spanbert-base','spanbert-large','luke-base','luke-large',
+'simcse-bert-base','simcse-bert-large','simcse-roberta-base','simcse-roberta-large']
+
+MAP_TYPES = ['usgs']
+CANDSET_MODES = ['all_map'] # candidate set is constructed based on all maps or one map
+
+
+def disambiguify(model, model_name, usgs_dataset, wikidata_dict_list, candset_mode = 'all_map', if_use_distance = True, select_indices = None): 
+
+    if select_indices is None: 
+        select_indices = range(0, len(usgs_dataset))
+
+
+    assert(candset_mode in ['all_map','per_map'])
+
+    wikidata_emb_list = [wikidata_dict['wikidata_emb_list'] for wikidata_dict in wikidata_dict_list]
+    wikidata_uri_list = [wikidata_dict['wikidata_uri_list'] for wikidata_dict in wikidata_dict_list]
+    #wikidata_des_list = [wikidata_dict['wikidata_des_list'] for wikidata_dict in wikidata_dict_list]
+
+    if candset_mode == 'all_map':
+        wikidata_emb_list = [item for sublist in wikidata_emb_list for item in sublist] # flatten
+        wikidata_uri_list = [item for sublist in wikidata_uri_list for item in sublist] # flatten
+        #wikidata_des_list = [item for sublist in wikidata_des_list for item in sublist] # flatten
+
+    
+    ret_list = []
+    for i in select_indices:
+
+        if candset_mode == 'per_map':
+            usgs_entity = usgs_dataset[i]
+            wikidata_emb_list = wikidata_emb_list[i]
+            wikidata_uri_list = wikidata_uri_list[i]
+            #wikidata_des_list = wikidata_des_list[i]
+
+        elif candset_mode == 'all_map':
+            usgs_entity = usgs_dataset[i]
+        else:
+            raise NotImplementedError
+        
+        if model_name == 'spatial_bert-base' or model_name == 'spatial_bert-large':
+            usgs_emb = get_spatialbert_embedding(usgs_entity, model, use_distance = if_use_distance)
+        else:
+            usgs_emb = get_bert_embedding(usgs_entity, model)
+
+      
+        sim_matrix = 1 - sp.distance.cdist(np.array(wikidata_emb_list), np.array([usgs_emb]), 'cosine')
+        
+        closest_match_uri = sort_ref_closest_match(sim_matrix, wikidata_uri_list)
+        #closest_match_des = sort_ref_closest_match(sim_matrix, wikidata_des_list)
+
+            
+        sorted_sim_matrix = np.sort(sim_matrix, axis = 0)[::-1] # descending order
+
+        ret_dict = dict()
+        ret_dict['pivot_name'] = usgs_entity['pivot_name']
+        ret_dict['sorted_match_uri'] = [a[0] for a in closest_match_uri]
+        #ret_dict['sorted_match_des'] = [a[0] for a in closest_match_des]
+        ret_dict['sorted_sim_matrix'] = [a[0] for a in sorted_sim_matrix]
+
+        ret_list.append(ret_dict)
+
+    return ret_list 
+
+def entity_linking_func(args):
+
+    model_name = args.model_name
+    map_type = args.map_type 
+    candset_mode = args.candset_mode
+
+    usgs_distance_norm_factor = args.usgs_distance_norm_factor
+    spatial_dist_fill= args.spatial_dist_fill
+    sep_between_neighbors = args.sep_between_neighbors
+
+    spatial_bert_weight_dir = args.spatial_bert_weight_dir
+    spatial_bert_weight_name = args.spatial_bert_weight_name
+
+    if_no_spatial_distance = args.no_spatial_distance
+
+
+    assert model_name in MODEL_OPTIONS
+    assert map_type in MAP_TYPES
+    assert candset_mode in CANDSET_MODES
+
+
+    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+
+    if args.out_dir is None:
+
+        if model_name == 'spatial_bert-base' or model_name == 'spatial_bert-large':
+
+            if sep_between_neighbors:
+                spatialbert_output_dir_str = 'dnorm' + str(usgs_distance_norm_factor ) + '_distfill' + str(spatial_dist_fill) + '_sep' 
+            else:
+                spatialbert_output_dir_str = 'dnorm' + str(usgs_distance_norm_factor ) + '_distfill' + str(spatial_dist_fill) + '_nosep' 
+
+
+            checkpoint_ep = spatial_bert_weight_name.split('_')[3]
+            checkpoint_iter = spatial_bert_weight_name.split('_')[4]
+            loss_val = spatial_bert_weight_name.split('_')[5][:-4]
+
+            if if_no_spatial_distance:
+                linking_prediction_dir = 'linking_prediction_dir/abalation_no_distance/'
+            else:
+                linking_prediction_dir = 'linking_prediction_dir'
+
+            if model_name == 'spatial_bert-base':
+                out_dir  = os.path.join('/data2/zekun/', linking_prediction_dir,  spatialbert_output_dir_str) + '/' + map_type + '-' + model_name + '-' + checkpoint_ep + '-' + checkpoint_iter + '-' + loss_val
+            elif model_name == 'spatial_bert-large':
+                
+                freeze_str = spatial_bert_weight_dir.split('/')[-2].split('_')[1] # either 'freeze' or 'nofreeze'
+                out_dir  = os.path.join('/data2/zekun/', linking_prediction_dir,  spatialbert_output_dir_str) + '/' + map_type + '-' + model_name + '-' + checkpoint_ep + '-' + checkpoint_iter + '-' + loss_val + '-' + freeze_str
+
+            
+
+        else:
+            out_dir  = '/data2/zekun/baseline_linking_prediction_dir/'  + map_type + '-' + model_name 
+
+    else:
+        out_dir = args.out_dir
+
+    print('out_dir', out_dir)
+
+    if not os.path.isdir(out_dir):
+        os.makedirs(out_dir)
+    
+        
+    if model_name == 'spatial_bert-base':
+        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+        
+        config = SpatialBertConfig()
+        model = SpatialBertModel(config)
+
+        model.to(device)
+        model.eval()
+        
+        # load pretrained weights
+        weight_path = os.path.join(spatial_bert_weight_dir, spatial_bert_weight_name)
+        model = load_spatial_bert_pretrained_weights(model, weight_path)
+
+    elif model_name == 'spatial_bert-large':
+        tokenizer = BertTokenizer.from_pretrained('bert-large-uncased')
+        
+        config = SpatialBertConfig(hidden_size = 1024, intermediate_size = 4096, num_attention_heads=16, num_hidden_layers=24)
+        model = SpatialBertModel(config)
+
+        model.to(device)
+        model.eval()
+        
+        # load pretrained weights
+        weight_path = os.path.join(spatial_bert_weight_dir, spatial_bert_weight_name)
+        model = load_spatial_bert_pretrained_weights(model, weight_path)
+
+    else:
+        model, tokenizer = get_baseline_model(model_name)
+        model.to(device)
+        model.eval()
+
+
+    if map_type == 'usgs':
+        map_name_list = ['USGS-15-CA-brawley-e1957-s1957-p1961',
+                        'USGS-15-CA-paloalto-e1899-s1895-rp1911',
+                        'USGS-15-CA-capesanmartin-e1921-s1917',
+                        'USGS-15-CA-sanfrancisco-e1899-s1892-rp1911',
+                        'USGS-30-CA-dardanelles-e1898-s1891-rp1912',
+                        'USGS-30-CA-holtville-e1907-s1905-rp1946',
+                        'USGS-30-CA-indiospecial-e1904-s1901-rp1910',
+                        'USGS-30-CA-lompoc-e1943-s1903-ap1941-rv1941',
+                        'USGS-30-CA-sanpedro-e1943-rv1944',
+                        'USGS-60-CA-alturas-e1892-rp1904',
+                        'USGS-60-CA-amboy-e1942',
+                        'USGS-60-CA-amboy-e1943-rv1943',
+                        'USGS-60-CA-modoclavabed-e1886-s1884',
+                        'USGS-60-CA-saltonsea-e1943-ap1940-rv1942']
+        
+    print('processing wikidata...')
+
+    wikidata_dict_list = []
+
+    wikidata_random30k = Wikidata_Random_Dataset(
+            data_file_path = '../data_processing/wikidata_sample30k/wikidata_30k_neighbor_reformat.json',
+            #neighbor_file_path = '../data_processing/wikidata_sample30k/wikidata_30k_neighbor.json',
+            tokenizer = tokenizer,
+            max_token_len = 512, 
+            distance_norm_factor = 0.0001, 
+            spatial_dist_fill=100, 
+            sep_between_neighbors = sep_between_neighbors,
+        )
+    
+    # process candidates for each phrase
+    for wikidata_cand in wikidata_random30k:
+        if model_name == 'spatial_bert-base' or model_name ==  'spatial_bert-large':
+            wikidata_emb = get_spatialbert_embedding(wikidata_cand, model)
+        else:
+            wikidata_emb = get_bert_embedding(wikidata_cand, model)
+            
+        wikidata_dict = {}
+        wikidata_dict['wikidata_emb_list'] = [wikidata_emb]
+        wikidata_dict['wikidata_uri_list'] = [wikidata_cand['uri']]
+
+        wikidata_dict_list.append(wikidata_dict)
+
+
+
+    for map_name in map_name_list:
+        
+        print(map_name)
+
+        wikidata_dict_per_map = {}
+        wikidata_dict_per_map['wikidata_emb_list'] = []
+        wikidata_dict_per_map['wikidata_uri_list'] = []
+
+        wikidata_dataset_permap =  Wikidata_Geocoord_Dataset(
+            data_file_path = '../data_processing/outputs/wikidata_reformat/wikidata_' + map_name + '.json',
+            tokenizer = tokenizer,
+            max_token_len = 512, 
+            distance_norm_factor = 0.0001, 
+            spatial_dist_fill=100,
+            sep_between_neighbors = sep_between_neighbors)
+
+
+
+        for i in range(0, len(wikidata_dataset_permap)):
+            # get all candiates for phrases within the map
+            wikidata_candidates = wikidata_dataset_permap[i] # dataset for each map, list of [cand for each phrase] 
+
+
+            # process candidates for each phrase
+            for wikidata_cand in wikidata_candidates:
+                if model_name == 'spatial_bert-base' or model_name ==  'spatial_bert-large':
+                    wikidata_emb = get_spatialbert_embedding(wikidata_cand, model)
+                else:
+                    wikidata_emb = get_bert_embedding(wikidata_cand, model)
+
+                wikidata_dict_per_map['wikidata_emb_list'].append(wikidata_emb)
+                wikidata_dict_per_map['wikidata_uri_list'].append(wikidata_cand['uri'])
+
+            
+        wikidata_dict_list.append(wikidata_dict_per_map)
+        
+
+
+    for map_name in map_name_list:
+        
+        print(map_name)
+        
+
+        usgs_dataset =  USGS_MapDataset(
+            data_file_path = '../data_processing/outputs/alignment_dir/map_' + map_name + '.json',
+            tokenizer = tokenizer,
+            distance_norm_factor = usgs_distance_norm_factor,
+            spatial_dist_fill = spatial_dist_fill,
+            sep_between_neighbors = sep_between_neighbors)
+        
+        
+        ret_list = disambiguify(model, model_name, usgs_dataset, wikidata_dict_list, candset_mode= candset_mode, if_use_distance = not if_no_spatial_distance, select_indices = None)
+
+        write_to_csv(out_dir, map_name, ret_list)
+
+    print('Done')
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--model_name', type=str, default='spatial_bert-base')
+    parser.add_argument('--out_dir', type=str, default=None)
+    parser.add_argument('--map_type', type=str, default='usgs')
+    parser.add_argument('--candset_mode', type=str, default='all_map')
+
+    parser.add_argument('--usgs_distance_norm_factor', type=float, default = 1)
+    parser.add_argument('--spatial_dist_fill', type=float, default = 100)
+                       
+    parser.add_argument('--sep_between_neighbors', default=False, action='store_true')
+    parser.add_argument('--no_spatial_distance',  default=False, action='store_true')
+
+    parser.add_argument('--spatial_bert_weight_dir', type = str, default = None)
+    parser.add_argument('--spatial_bert_weight_name', type = str, default = None)
+                        
+    args = parser.parse_args()
+    print('\n')
+    print(args)
+    print('\n')
+
+    # out_dir not None, and out_dir does not exist, then create out_dir
+    if args.out_dir is not None and not os.path.isdir(args.out_dir):
+        os.makedirs(args.out_dir)
+
+    entity_linking_func(args)
+
+
+
+if __name__ == '__main__':
+
+    main()
+
+    

models/spabert/experiments/semantic_typing/data_processing/merge_osm_json.py

@@ -0,0 +1,97 @@
+import os
+import json 
+import math
+import glob
+import re
+import pdb
+
+'''
+NO LONGER NEEDED
+
+Process the california, london, and minnesota OSM data and prepare pseudo-sentence, spatial context
+
+Load the raw output files genrated by sql
+Unify the json by changing the structure of dictionary
+Save the output into two files, one for training+ validation, and the other one for testing
+'''
+
+region_list = ['california','london','minnesota']
+
+input_json_dir = '../data/sql_output/sub_files/'
+output_json_dir = '../data/sql_output/'
+
+for region_name in region_list:
+    file_list = glob.glob(os.path.join(input_json_dir, 'spatialbert-osm-point-' + region_name + '*.json'))
+    file_list = sorted(file_list)
+    print('found %d files for region %s' % (len(file_list), region_name))
+
+    
+    num_test_files = int(math.ceil(len(file_list) * 0.2))
+    num_train_val_files = len(file_list) - num_test_files
+
+    print('%d files for train-val' % num_train_val_files)
+    print('%d files for test-tes' % num_test_files)
+
+    train_val_output_path = os.path.join(output_json_dir + 'osm-point-' + region_name + '_train_val.json')
+    test_output_path = os.path.join(output_json_dir + 'osm-point-' + region_name + '_test.json')
+
+    # refresh the file
+    with open(train_val_output_path, 'w') as f:
+        pass
+    with open(test_output_path, 'w') as f:
+        pass
+
+    for idx in range(len(file_list)):
+
+        if idx < num_train_val_files:
+            output_path = train_val_output_path
+        else:
+            output_path = test_output_path
+
+        file_path = file_list[idx]
+
+        print(file_path)
+        
+        with open(file_path, 'r') as f:
+            data = f.readlines()
+
+        line = data[0]
+
+        line = re.sub(r'\n', '', line)
+        line = re.sub(r'\\n', '', line)
+        line = re.sub(r'\\+', '', line)
+        line = re.sub(r'\+', '', line)
+
+        line_dict_list = json.loads(line)
+
+
+        for line_dict in line_dict_list:
+            
+            line_dict = line_dict['json_build_object']
+
+            if not line_dict['name'][0].isalpha(): # discard record if the first char is not enghlish etter
+                continue
+
+            neighbor_name_list = line_dict['neighbor_info'][0]['name_list']
+            neighbor_geom_list = line_dict['neighbor_info'][0]['geometry_list']
+
+            assert(len(neighbor_geom_list) == len(neighbor_geom_list))
+
+            temp_dict = \
+            {'info':{'name':line_dict['name'], 
+                    'geometry':{'coordinates':line_dict['geometry']},
+                    'class':line_dict['class']
+                    }, 
+            'neighbor_info':{'name_list': neighbor_name_list,
+                            'geometry_list': neighbor_geom_list
+                    }
+            }
+
+            with open(output_path, 'a') as f:
+                json.dump(temp_dict, f)
+                f.write('\n')
+
+        
+
+            
+                        

models/spabert/experiments/semantic_typing/src/run_baseline_test.py

@@ -0,0 +1,82 @@
+import os 
+import pdb
+import argparse
+import numpy as np
+import time
+
+MODEL_OPTIONS = ['bert-base','bert-large','roberta-base','roberta-large',
+'spanbert-base','spanbert-large','luke-base','luke-large',
+'simcse-bert-base','simcse-bert-large','simcse-roberta-base','simcse-roberta-large']
+
+def execute_command(command, if_print_command):
+    t1 = time.time()
+
+    if if_print_command:
+        print(command)
+    os.system(command)
+
+    t2 = time.time()
+    time_usage = t2 - t1 
+    return time_usage
+
+def run_test(args):
+    weight_dir = args.weight_dir
+    backbone_option = args.backbone_option
+    gpu_id = str(args.gpu_id)
+    if_print_command = args.print_command
+    sep_between_neighbors =  args.sep_between_neighbors
+
+    assert backbone_option in MODEL_OPTIONS
+
+    if sep_between_neighbors:
+        sep_str = '_sep'
+    else:
+        sep_str = ''
+    
+    if 'large' in backbone_option:
+        checkpoint_dir = os.path.join(weight_dir, 'typing_lr1e-06_%s_nofreeze%s_london_california_bsize12'% (backbone_option, sep_str))
+    else:
+        checkpoint_dir = os.path.join(weight_dir, 'typing_lr5e-05_%s_nofreeze%s_london_california_bsize12'% (backbone_option, sep_str))
+    weight_files = os.listdir(checkpoint_dir)
+
+    val_loss_list = [weight_file.split('_')[-1] for weight_file in weight_files]
+    min_loss_weight = weight_files[np.argmin(val_loss_list)]
+
+    checkpoint_path = os.path.join(checkpoint_dir, min_loss_weight)
+
+    if sep_between_neighbors:
+        command = 'CUDA_VISIBLE_DEVICES=%s python3 test_cls_baseline.py  --sep_between_neighbors --backbone_option=%s --batch_size=8 --with_type --checkpoint_path=%s ' % (gpu_id, backbone_option, checkpoint_path)
+    else:
+        command = 'CUDA_VISIBLE_DEVICES=%s python3 test_cls_baseline.py  --backbone_option=%s --batch_size=8 --with_type --checkpoint_path=%s ' % (gpu_id, backbone_option, checkpoint_path)
+
+
+    execute_command(command, if_print_command)
+
+
+
+
+def main():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument('--weight_dir', type=str, default='/data2/zekun/spatial_bert_baseline_weights/')
+    parser.add_argument('--sep_between_neighbors', default=False, action='store_true')
+    parser.add_argument('--backbone_option', type=str, default=None)
+    parser.add_argument('--gpu_id', type=int, default=0) # output prefix 
+
+    parser.add_argument('--print_command', default=False, action='store_true')
+
+                        
+    args = parser.parse_args()
+    print('\n')
+    print(args)
+    print('\n')
+
+    run_test(args)
+
+
+if __name__ == '__main__':
+
+    main()
+
+    
+

models/spabert/experiments/semantic_typing/src/test_cls_ablation_spatialbert.py

@@ -0,0 +1,209 @@
+import os
+import sys
+from transformers import RobertaTokenizer, BertTokenizer
+from tqdm import tqdm  # for our progress bar
+from transformers import AdamW
+
+import torch
+from torch.utils.data import DataLoader
+import torch.nn.functional as F
+
+sys.path.append('../../../')
+from models.spatial_bert_model import SpatialBertModel
+from models.spatial_bert_model import SpatialBertConfig
+from models.spatial_bert_model import  SpatialBertForMaskedLM, SpatialBertForSemanticTyping
+from datasets.osm_sample_loader import PbfMapDataset
+from datasets.const import *
+from transformers.models.bert.modeling_bert import BertForMaskedLM
+
+from sklearn.metrics import label_ranking_average_precision_score
+from sklearn.metrics import precision_recall_fscore_support
+import numpy as np
+import argparse 
+from sklearn.preprocessing import LabelEncoder
+import pdb
+
+
+DEBUG = False
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.benchmark = False
+torch.manual_seed(42)
+torch.cuda.manual_seed_all(42)
+
+
+def testing(args):
+
+    max_token_len = args.max_token_len
+    batch_size = args.batch_size
+    num_workers = args.num_workers
+    distance_norm_factor = args.distance_norm_factor
+    spatial_dist_fill=args.spatial_dist_fill
+    with_type = args.with_type
+    sep_between_neighbors = args.sep_between_neighbors
+    checkpoint_path = args.checkpoint_path
+    if_no_spatial_distance = args.no_spatial_distance
+
+    bert_option = args.bert_option
+    num_neighbor_limit = args.num_neighbor_limit
+
+
+    london_file_path = '../data/sql_output/osm-point-london-typing.json'
+    california_file_path = '../data/sql_output/osm-point-california-typing.json'
+
+    
+    if bert_option == 'bert-base':
+        tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
+        config = SpatialBertConfig(use_spatial_distance_embedding = not if_no_spatial_distance, num_semantic_types=len(CLASS_9_LIST))
+    elif bert_option == 'bert-large':
+        tokenizer = BertTokenizer.from_pretrained("bert-large-uncased")
+        config = SpatialBertConfig(use_spatial_distance_embedding = not if_no_spatial_distance, hidden_size = 1024, intermediate_size = 4096, num_attention_heads=16, num_hidden_layers=24,num_semantic_types=len(CLASS_9_LIST))
+    else:
+        raise NotImplementedError
+
+
+    model = SpatialBertForSemanticTyping(config)
+
+
+    label_encoder = LabelEncoder()
+    label_encoder.fit(CLASS_9_LIST)
+
+    
+
+    london_dataset = PbfMapDataset(data_file_path = london_file_path, 
+                                        tokenizer = tokenizer, 
+                                        max_token_len = max_token_len, 
+                                        distance_norm_factor = distance_norm_factor, 
+                                        spatial_dist_fill = spatial_dist_fill, 
+                                        with_type = with_type,
+                                        sep_between_neighbors = sep_between_neighbors, 
+                                        label_encoder = label_encoder,
+                                        num_neighbor_limit = num_neighbor_limit, 
+                                        mode = 'test',)
+
+    california_dataset = PbfMapDataset(data_file_path = california_file_path, 
+                                            tokenizer = tokenizer, 
+                                            max_token_len = max_token_len, 
+                                            distance_norm_factor = distance_norm_factor, 
+                                            spatial_dist_fill = spatial_dist_fill, 
+                                            with_type = with_type,
+                                            sep_between_neighbors = sep_between_neighbors,
+                                            label_encoder = label_encoder,
+                                            num_neighbor_limit = num_neighbor_limit,
+                                            mode = 'test')
+
+    test_dataset = torch.utils.data.ConcatDataset([london_dataset, california_dataset])
+
+
+    
+    test_loader = DataLoader(test_dataset, batch_size= batch_size, num_workers=num_workers,
+                            shuffle=False, pin_memory=True, drop_last=False)
+    
+
+    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+    model.to(device)
+    
+
+    model.load_state_dict(torch.load(checkpoint_path)) # #
+
+    model.eval()
+
+
+
+    print('start testing...')
+
+    
+    # setup loop with TQDM and dataloader
+    loop = tqdm(test_loader, leave=True)
+
+
+    mrr_total = 0.
+    prec_total = 0.
+    sample_cnt = 0
+
+    gt_list = []
+    pred_list = []
+
+    for batch in loop:
+        # initialize calculated gradients (from prev step)
+        
+        # pull all tensor batches required for training
+        input_ids = batch['pseudo_sentence'].to(device)
+        attention_mask = batch['attention_mask'].to(device)
+        position_list_x = batch['norm_lng_list'].to(device)
+        position_list_y = batch['norm_lat_list'].to(device)
+        sent_position_ids = batch['sent_position_ids'].to(device)
+
+        #labels = batch['pseudo_sentence'].to(device)
+        labels = batch['pivot_type'].to(device)
+        pivot_lens = batch['pivot_token_len'].to(device)
+
+        outputs = model(input_ids, attention_mask = attention_mask, sent_position_ids = sent_position_ids,
+            position_list_x = position_list_x, position_list_y = position_list_y, labels = labels, pivot_len_list = pivot_lens)
+
+
+        onehot_labels = F.one_hot(labels, num_classes=9)
+        
+        gt_list.extend(onehot_labels.cpu().detach().numpy())
+        pred_list.extend(outputs.logits.cpu().detach().numpy())
+
+        #pdb.set_trace()
+        mrr = label_ranking_average_precision_score(onehot_labels.cpu().detach().numpy(), outputs.logits.cpu().detach().numpy())
+        mrr_total += mrr * input_ids.shape[0] 
+        sample_cnt += input_ids.shape[0]
+
+    precisions, recalls, fscores, supports = precision_recall_fscore_support(np.argmax(np.array(gt_list),axis=1), np.argmax(np.array(pred_list),axis=1), average=None)
+
+    precision, recall, f1, _ = precision_recall_fscore_support(np.argmax(np.array(gt_list),axis=1), np.argmax(np.array(pred_list),axis=1), average='micro')
+ 
+    # print('precisions:\n', ["{:.3f}".format(prec) for prec in precisions])
+    # print('recalls:\n', ["{:.3f}".format(rec) for rec in recalls])
+    # print('fscores:\n', ["{:.3f}".format(f1) for f1 in fscores])
+    # print('supports:\n', supports)
+    print('micro P, micro R, micro F1', "{:.3f}".format(precision), "{:.3f}".format(recall), "{:.3f}".format(f1))
+    
+
+
+def main():
+
+    parser = argparse.ArgumentParser()
+    
+    parser.add_argument('--max_token_len', type=int, default=300)
+    parser.add_argument('--batch_size', type=int, default=12)
+    parser.add_argument('--num_workers', type=int, default=5)
+
+    parser.add_argument('--num_neighbor_limit', type=int, default = None)
+
+    parser.add_argument('--distance_norm_factor', type=float, default = 0.0001)
+    parser.add_argument('--spatial_dist_fill', type=float, default = 20)
+
+    
+    parser.add_argument('--with_type', default=False, action='store_true')
+    parser.add_argument('--sep_between_neighbors', default=False, action='store_true')
+    parser.add_argument('--no_spatial_distance',  default=False, action='store_true')
+
+    parser.add_argument('--bert_option', type=str, default='bert-base')
+    parser.add_argument('--prediction_save_dir', type=str, default=None)
+
+    parser.add_argument('--checkpoint_path', type=str, default=None)
+
+
+                        
+    args = parser.parse_args()
+    print('\n')
+    print(args)
+    print('\n')
+
+
+    # out_dir not None, and out_dir does not exist, then create out_dir
+    if args.prediction_save_dir is not None and not os.path.isdir(args.prediction_save_dir):
+        os.makedirs(args.prediction_save_dir)
+
+    testing(args)
+
+   
+
+if __name__ == '__main__':
+
+    main()
+
+    

models/spabert/experiments/semantic_typing/src/test_cls_baseline.py

@@ -0,0 +1,189 @@
+import os
+import sys
+from tqdm import tqdm  # for our progress bar
+import numpy as np
+import argparse 
+from sklearn.preprocessing import LabelEncoder
+import pdb
+
+
+import torch
+from torch.utils.data import DataLoader
+from transformers import AdamW
+import torch.nn.functional as F
+
+sys.path.append('../../../')
+from datasets.osm_sample_loader import PbfMapDataset
+from datasets.const import *
+from utils.baseline_utils import get_baseline_model
+from models.baseline_typing_model import BaselineForSemanticTyping
+
+from sklearn.metrics import label_ranking_average_precision_score
+from sklearn.metrics import precision_recall_fscore_support
+
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.benchmark = False
+torch.manual_seed(42)
+torch.cuda.manual_seed_all(42)
+
+
+
+MODEL_OPTIONS = ['bert-base','bert-large','roberta-base','roberta-large',
+'spanbert-base','spanbert-large','luke-base','luke-large',
+'simcse-bert-base','simcse-bert-large','simcse-roberta-base','simcse-roberta-large']
+
+def testing(args):
+
+    num_workers = args.num_workers
+    batch_size = args.batch_size
+    max_token_len = args.max_token_len
+
+    distance_norm_factor = args.distance_norm_factor
+    spatial_dist_fill=args.spatial_dist_fill
+    with_type = args.with_type
+    sep_between_neighbors = args.sep_between_neighbors
+    freeze_backbone = args.freeze_backbone
+
+
+    backbone_option = args.backbone_option
+
+    checkpoint_path = args.checkpoint_path
+
+    assert(backbone_option in MODEL_OPTIONS)
+    
+   
+    london_file_path = '../data/sql_output/osm-point-london-typing.json'
+    california_file_path = '../data/sql_output/osm-point-california-typing.json'
+
+
+    
+    label_encoder = LabelEncoder()
+    label_encoder.fit(CLASS_9_LIST)
+
+
+    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+
+    backbone_model, tokenizer = get_baseline_model(backbone_option)
+    model = BaselineForSemanticTyping(backbone_model, backbone_model.config.hidden_size, len(CLASS_9_LIST))
+
+    model.load_state_dict(torch.load(checkpoint_path) ) #, strict = False # load sentence position embedding weights as well
+
+    model.to(device)
+    model.train()
+
+
+
+    london_dataset = PbfMapDataset(data_file_path = london_file_path, 
+                                        tokenizer = tokenizer, 
+                                        max_token_len = max_token_len, 
+                                        distance_norm_factor = distance_norm_factor, 
+                                        spatial_dist_fill = spatial_dist_fill, 
+                                        with_type = with_type,
+                                        sep_between_neighbors = sep_between_neighbors, 
+                                        label_encoder = label_encoder,
+                                        mode = 'test')
+
+    
+    california_dataset = PbfMapDataset(data_file_path = california_file_path, 
+                                            tokenizer = tokenizer, 
+                                            max_token_len = max_token_len, 
+                                            distance_norm_factor = distance_norm_factor, 
+                                            spatial_dist_fill = spatial_dist_fill, 
+                                            with_type = with_type,
+                                            sep_between_neighbors = sep_between_neighbors,
+                                            label_encoder = label_encoder,
+                                            mode = 'test')
+
+    test_dataset = torch.utils.data.ConcatDataset([london_dataset, california_dataset])
+
+    
+    test_loader = DataLoader(test_dataset, batch_size= batch_size, num_workers=num_workers,
+                                shuffle=False, pin_memory=True, drop_last=False)
+    
+
+    
+
+   
+    print('start testing...')
+
+    # setup loop with TQDM and dataloader
+    loop = tqdm(test_loader, leave=True)
+
+    mrr_total = 0.
+    prec_total = 0.
+    sample_cnt = 0
+
+    gt_list = []
+    pred_list = []
+
+    for batch in loop:
+        # initialize calculated gradients (from prev step)
+        
+        # pull all tensor batches required for training
+        input_ids = batch['pseudo_sentence'].to(device)
+        attention_mask = batch['attention_mask'].to(device)
+        position_ids = batch['sent_position_ids'].to(device)
+
+        #labels = batch['pseudo_sentence'].to(device)
+        labels = batch['pivot_type'].to(device)
+        pivot_lens = batch['pivot_token_len'].to(device)
+
+        outputs = model(input_ids, attention_mask = attention_mask, position_ids = position_ids,
+                labels = labels, pivot_len_list = pivot_lens)
+
+
+        onehot_labels = F.one_hot(labels, num_classes=9)
+        
+        gt_list.extend(onehot_labels.cpu().detach().numpy())
+        pred_list.extend(outputs.logits.cpu().detach().numpy())
+
+        mrr = label_ranking_average_precision_score(onehot_labels.cpu().detach().numpy(), outputs.logits.cpu().detach().numpy())
+        mrr_total += mrr * input_ids.shape[0] 
+        sample_cnt += input_ids.shape[0]
+
+    precisions, recalls, fscores, supports = precision_recall_fscore_support(np.argmax(np.array(gt_list),axis=1), np.argmax(np.array(pred_list),axis=1), average=None)
+    precision, recall, f1, _ = precision_recall_fscore_support(np.argmax(np.array(gt_list),axis=1), np.argmax(np.array(pred_list),axis=1), average='micro')
+    print('precisions:\n', ["{:.3f}".format(prec) for prec in precisions])
+    print('recalls:\n', ["{:.3f}".format(rec) for rec in recalls])
+    print('fscores:\n', ["{:.3f}".format(f1) for f1 in fscores])
+    print('supports:\n', supports)
+    print('micro P, micro R, micro F1', "{:.3f}".format(precision), "{:.3f}".format(recall), "{:.3f}".format(f1))
+    
+    #pdb.set_trace()
+    #print(mrr_total/sample_cnt)
+
+
+       
+def main():
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--num_workers', type=int, default=5)
+    parser.add_argument('--batch_size', type=int, default=12)
+    parser.add_argument('--max_token_len', type=int, default=300)
+    
+
+    parser.add_argument('--distance_norm_factor', type=float, default = 0.0001)
+    parser.add_argument('--spatial_dist_fill', type=float, default = 20)
+
+    parser.add_argument('--with_type', default=False, action='store_true')
+    parser.add_argument('--sep_between_neighbors', default=False, action='store_true')
+    parser.add_argument('--freeze_backbone', default=False, action='store_true')
+
+    parser.add_argument('--backbone_option', type=str, default='bert-base')
+    parser.add_argument('--checkpoint_path', type=str, default=None)
+
+                        
+    args = parser.parse_args()
+    print('\n')
+    print(args)
+    print('\n')
+
+
+    testing(args)
+
+
+if __name__ == '__main__':
+
+    main()
+
+    

models/spabert/experiments/semantic_typing/src/test_cls_spatialbert.py

@@ -0,0 +1,214 @@
+import os
+import sys
+from transformers import RobertaTokenizer, BertTokenizer
+from tqdm import tqdm  # for our progress bar
+from transformers import AdamW
+
+import torch
+from torch.utils.data import DataLoader
+import torch.nn.functional as F
+
+sys.path.append('../../../')
+from models.spatial_bert_model import SpatialBertModel
+from models.spatial_bert_model import SpatialBertConfig
+from models.spatial_bert_model import  SpatialBertForMaskedLM, SpatialBertForSemanticTyping
+from datasets.osm_sample_loader import PbfMapDataset
+from datasets.const import *
+from transformers.models.bert.modeling_bert import BertForMaskedLM
+
+from sklearn.metrics import label_ranking_average_precision_score
+from sklearn.metrics import precision_recall_fscore_support
+import numpy as np
+import argparse 
+from sklearn.preprocessing import LabelEncoder
+import pdb
+
+
+DEBUG = False
+
+torch.backends.cudnn.deterministic = True
+torch.backends.cudnn.benchmark = False
+torch.manual_seed(42)
+torch.cuda.manual_seed_all(42)
+
+
+def testing(args):
+
+    max_token_len = args.max_token_len
+    batch_size = args.batch_size
+    num_workers = args.num_workers
+    distance_norm_factor = args.distance_norm_factor
+    spatial_dist_fill=args.spatial_dist_fill
+    with_type = args.with_type
+    sep_between_neighbors = args.sep_between_neighbors
+    checkpoint_path = args.checkpoint_path
+    if_no_spatial_distance = args.no_spatial_distance
+
+    bert_option = args.bert_option
+
+
+
+    if args.num_classes == 9:
+        london_file_path = '../../semantic_typing/data/sql_output/osm-point-london-typing.json'
+        california_file_path = '../../semantic_typing/data/sql_output/osm-point-california-typing.json'
+        TYPE_LIST = CLASS_9_LIST
+        type_key_str  = 'class'
+    elif args.num_classes == 74:
+        london_file_path = '../../semantic_typing/data/sql_output/osm-point-london-typing-ranking.json'
+        california_file_path = '../../semantic_typing/data/sql_output/osm-point-california-typing-ranking.json'
+        TYPE_LIST = CLASS_74_LIST
+        type_key_str = 'fine_class'
+    else:
+        raise NotImplementedError
+
+    if bert_option == 'bert-base':
+        tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
+        config = SpatialBertConfig(use_spatial_distance_embedding = not if_no_spatial_distance, num_semantic_types=len(TYPE_LIST))
+    elif bert_option == 'bert-large':
+        tokenizer = BertTokenizer.from_pretrained("bert-large-uncased")
+        config = SpatialBertConfig(use_spatial_distance_embedding = not if_no_spatial_distance, hidden_size = 1024, intermediate_size = 4096, num_attention_heads=16, num_hidden_layers=24,num_semantic_types=len(TYPE_LIST))
+    else:
+        raise NotImplementedError
+
+
+    model = SpatialBertForSemanticTyping(config)
+
+
+    label_encoder = LabelEncoder()
+    label_encoder.fit(TYPE_LIST)
+
+    london_dataset = PbfMapDataset(data_file_path = london_file_path, 
+                                        tokenizer = tokenizer, 
+                                        max_token_len = max_token_len, 
+                                        distance_norm_factor = distance_norm_factor, 
+                                        spatial_dist_fill = spatial_dist_fill, 
+                                        with_type = with_type,
+                                        type_key_str = type_key_str,
+                                        sep_between_neighbors = sep_between_neighbors, 
+                                        label_encoder = label_encoder,
+                                        mode = 'test')
+
+    california_dataset = PbfMapDataset(data_file_path = california_file_path, 
+                                            tokenizer = tokenizer, 
+                                            max_token_len = max_token_len, 
+                                            distance_norm_factor = distance_norm_factor, 
+                                            spatial_dist_fill = spatial_dist_fill, 
+                                            with_type = with_type,
+                                            type_key_str = type_key_str,
+                                            sep_between_neighbors = sep_between_neighbors,
+                                            label_encoder = label_encoder,
+                                            mode = 'test')
+
+    test_dataset = torch.utils.data.ConcatDataset([london_dataset, california_dataset])
+
+
+    
+    test_loader = DataLoader(test_dataset, batch_size= batch_size, num_workers=num_workers,
+                            shuffle=False, pin_memory=True, drop_last=False)
+    
+
+    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+    model.to(device)
+
+    model.load_state_dict(torch.load(checkpoint_path))
+
+    model.eval()
+
+
+
+    print('start testing...')
+
+    
+    # setup loop with TQDM and dataloader
+    loop = tqdm(test_loader, leave=True)
+
+
+    mrr_total = 0.
+    prec_total = 0.
+    sample_cnt = 0
+
+    gt_list = []
+    pred_list = []
+
+    for batch in loop:
+        # initialize calculated gradients (from prev step)
+        
+        # pull all tensor batches required for training
+        input_ids = batch['pseudo_sentence'].to(device)
+        attention_mask = batch['attention_mask'].to(device)
+        position_list_x = batch['norm_lng_list'].to(device)
+        position_list_y = batch['norm_lat_list'].to(device)
+        sent_position_ids = batch['sent_position_ids'].to(device)
+
+        #labels = batch['pseudo_sentence'].to(device)
+        labels = batch['pivot_type'].to(device)
+        pivot_lens = batch['pivot_token_len'].to(device)
+
+        outputs = model(input_ids, attention_mask = attention_mask, sent_position_ids = sent_position_ids,
+            position_list_x = position_list_x, position_list_y = position_list_y, labels = labels, pivot_len_list = pivot_lens)
+
+
+        onehot_labels = F.one_hot(labels, num_classes=len(TYPE_LIST))
+        
+        gt_list.extend(onehot_labels.cpu().detach().numpy())
+        pred_list.extend(outputs.logits.cpu().detach().numpy())
+
+        mrr = label_ranking_average_precision_score(onehot_labels.cpu().detach().numpy(), outputs.logits.cpu().detach().numpy())
+        mrr_total += mrr * input_ids.shape[0] 
+        sample_cnt += input_ids.shape[0]
+
+    precisions, recalls, fscores, supports = precision_recall_fscore_support(np.argmax(np.array(gt_list),axis=1), np.argmax(np.array(pred_list),axis=1), average=None)
+    # print('precisions:\n', precisions)
+    # print('recalls:\n', recalls)
+    # print('fscores:\n', fscores)
+    # print('supports:\n', supports)
+    precision, recall, f1, _ = precision_recall_fscore_support(np.argmax(np.array(gt_list),axis=1), np.argmax(np.array(pred_list),axis=1), average='micro')
+    print('precisions:\n', ["{:.3f}".format(prec) for prec in precisions])
+    print('recalls:\n', ["{:.3f}".format(rec) for rec in recalls])
+    print('fscores:\n', ["{:.3f}".format(f1) for f1 in fscores])
+    print('supports:\n', supports)
+    print('micro P, micro R, micro F1', "{:.3f}".format(precision), "{:.3f}".format(recall), "{:.3f}".format(f1))
+    
+
+
+
+def main():
+
+    parser = argparse.ArgumentParser()
+    
+    parser.add_argument('--max_token_len', type=int, default=512)
+    parser.add_argument('--batch_size', type=int, default=12)
+    parser.add_argument('--num_workers', type=int, default=5)
+
+    parser.add_argument('--distance_norm_factor', type=float, default = 0.0001)
+    parser.add_argument('--spatial_dist_fill', type=float, default = 100)
+    parser.add_argument('--num_classes', type=int, default = 9)
+
+    parser.add_argument('--with_type', default=False, action='store_true')
+    parser.add_argument('--sep_between_neighbors', default=False, action='store_true')
+    parser.add_argument('--no_spatial_distance',  default=False, action='store_true')
+
+    parser.add_argument('--bert_option', type=str, default='bert-base')
+    parser.add_argument('--prediction_save_dir', type=str, default=None)
+
+    parser.add_argument('--checkpoint_path', type=str, default=None)
+
+                        
+    args = parser.parse_args()
+    print('\n')
+    print(args)
+    print('\n')
+
+
+    # out_dir not None, and out_dir does not exist, then create out_dir
+    if args.prediction_save_dir is not None and not os.path.isdir(args.prediction_save_dir):
+        os.makedirs(args.prediction_save_dir)
+
+    testing(args)
+
+  
+if __name__ == '__main__':
+
+    main()
+
+    

models/spabert/experiments/semantic_typing/src/train_cls_baseline.py

@@ -0,0 +1,227 @@
+import os
+import sys
+from tqdm import tqdm  # for our progress bar
+import numpy as np
+import argparse 
+from sklearn.preprocessing import LabelEncoder
+import pdb
+
+
+import torch
+from torch.utils.data import DataLoader
+from transformers import AdamW
+
+sys.path.append('../../../')
+from datasets.osm_sample_loader import PbfMapDataset
+from datasets.const import *
+from utils.baseline_utils import get_baseline_model
+from models.baseline_typing_model import BaselineForSemanticTyping
+
+
+MODEL_OPTIONS = ['bert-base','bert-large','roberta-base','roberta-large',
+'spanbert-base','spanbert-large','luke-base','luke-large',
+'simcse-bert-base','simcse-bert-large','simcse-roberta-base','simcse-roberta-large']
+
+def training(args):
+
+    num_workers = args.num_workers
+    batch_size = args.batch_size
+    epochs = args.epochs
+    lr = args.lr #1e-7 # 5e-5
+    save_interval = args.save_interval
+    max_token_len = args.max_token_len
+    distance_norm_factor = args.distance_norm_factor
+    spatial_dist_fill=args.spatial_dist_fill
+    with_type = args.with_type
+    sep_between_neighbors = args.sep_between_neighbors
+    freeze_backbone = args.freeze_backbone
+
+
+    backbone_option = args.backbone_option
+
+    assert(backbone_option in MODEL_OPTIONS)
+    
+   
+    london_file_path = '../data/sql_output/osm-point-london-typing.json'
+    california_file_path = '../data/sql_output/osm-point-california-typing.json'
+
+    if args.model_save_dir is None:
+        freeze_pathstr = '_freeze' if freeze_backbone else '_nofreeze'
+        sep_pathstr = '_sep' if sep_between_neighbors else '_nosep'
+        model_save_dir = '/data2/zekun/spatial_bert_baseline_weights/typing_lr' + str("{:.0e}".format(lr))  +'_'+backbone_option+ freeze_pathstr + sep_pathstr + '_london_california_bsize' + str(batch_size)
+        
+        if not os.path.isdir(model_save_dir):
+            os.makedirs(model_save_dir)
+    else:
+        model_save_dir = args.model_save_dir
+
+    print('model_save_dir', model_save_dir)
+    print('\n')
+
+    
+    label_encoder = LabelEncoder()
+    label_encoder.fit(CLASS_9_LIST)
+
+    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+
+    backbone_model, tokenizer = get_baseline_model(backbone_option)
+    model = BaselineForSemanticTyping(backbone_model, backbone_model.config.hidden_size, len(CLASS_9_LIST))
+
+    model.to(device)
+    model.train()
+
+    london_train_val_dataset = PbfMapDataset(data_file_path = london_file_path, 
+                                        tokenizer = tokenizer, 
+                                        max_token_len = max_token_len, 
+                                        distance_norm_factor = distance_norm_factor, 
+                                        spatial_dist_fill = spatial_dist_fill, 
+                                        with_type = with_type,
+                                        sep_between_neighbors = sep_between_neighbors, 
+                                        label_encoder = label_encoder,
+                                        mode = 'train')
+
+    percent_80 = int(len(london_train_val_dataset) * 0.8)
+    london_train_dataset, london_val_dataset = torch.utils.data.random_split(london_train_val_dataset, [percent_80, len(london_train_val_dataset) - percent_80])
+
+    california_train_val_dataset = PbfMapDataset(data_file_path = california_file_path, 
+                                            tokenizer = tokenizer, 
+                                            max_token_len = max_token_len, 
+                                            distance_norm_factor = distance_norm_factor, 
+                                            spatial_dist_fill = spatial_dist_fill, 
+                                            with_type = with_type,
+                                            sep_between_neighbors = sep_between_neighbors,
+                                            label_encoder = label_encoder,
+                                            mode = 'train')
+    percent_80 = int(len(california_train_val_dataset) * 0.8)
+    california_train_dataset, california_val_dataset = torch.utils.data.random_split(california_train_val_dataset, [percent_80, len(california_train_val_dataset) - percent_80])
+
+    train_dataset = torch.utils.data.ConcatDataset([london_train_dataset, california_train_dataset])
+    val_dataset = torch.utils.data.ConcatDataset([london_val_dataset, california_val_dataset])
+
+
+    
+    train_loader = DataLoader(train_dataset, batch_size= batch_size, num_workers=num_workers,
+                            shuffle=True, pin_memory=True, drop_last=True)
+    val_loader = DataLoader(val_dataset, batch_size= batch_size, num_workers=num_workers,
+                            shuffle=False, pin_memory=True, drop_last=False)
+
+
+
+    
+
+    
+
+    # initialize optimizer
+    optim = AdamW(model.parameters(), lr = lr)
+
+    print('start training...')
+
+    for epoch in range(epochs):
+        # setup loop with TQDM and dataloader
+        loop = tqdm(train_loader, leave=True)
+        iter = 0
+        for batch in loop:
+            # initialize calculated gradients (from prev step)
+            optim.zero_grad()
+            # pull all tensor batches required for training
+            input_ids = batch['pseudo_sentence'].to(device)
+            attention_mask = batch['attention_mask'].to(device)
+            position_ids = batch['sent_position_ids'].to(device)
+
+            #labels = batch['pseudo_sentence'].to(device)
+            labels = batch['pivot_type'].to(device)
+            pivot_lens = batch['pivot_token_len'].to(device)
+
+            outputs = model(input_ids, attention_mask = attention_mask, position_ids = position_ids,
+                labels = labels, pivot_len_list = pivot_lens)
+
+
+            loss = outputs.loss
+            loss.backward()
+            optim.step()
+
+            loop.set_description(f'Epoch {epoch}')
+            loop.set_postfix({'loss':loss.item()})
+            
+            
+            iter += 1
+
+            if iter % save_interval == 0 or iter == loop.total:
+                loss_valid = validating(val_loader, model, device)
+                print('validation loss', loss_valid)
+
+                save_path = os.path.join(model_save_dir, 'ep'+str(epoch) + '_iter'+ str(iter).zfill(5) \
+                + '_' +str("{:.4f}".format(loss.item())) + '_val' + str("{:.4f}".format(loss_valid)) +'.pth' )
+
+                torch.save(model.state_dict(), save_path)
+                print('saving model checkpoint to', save_path)
+
+                
+                
+def validating(val_loader, model, device):
+
+    with torch.no_grad():
+
+        loss_valid = 0
+        loop = tqdm(val_loader, leave=True)
+
+        for batch in loop:
+            input_ids = batch['pseudo_sentence'].to(device)
+            attention_mask = batch['attention_mask'].to(device)
+            position_ids = batch['sent_position_ids'].to(device)
+
+            labels = batch['pivot_type'].to(device)
+            pivot_lens = batch['pivot_token_len'].to(device)
+
+            outputs = model(input_ids, attention_mask = attention_mask, position_ids = position_ids,
+                 labels = labels, pivot_len_list = pivot_lens)
+
+            loss_valid  += outputs.loss
+
+        loss_valid /= len(val_loader)
+
+        return loss_valid
+
+
+def main():
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--num_workers', type=int, default=5)
+    parser.add_argument('--batch_size', type=int, default=12)
+    parser.add_argument('--epochs', type=int, default=10)
+    parser.add_argument('--save_interval', type=int, default=2000)
+    parser.add_argument('--max_token_len', type=int, default=300)
+    
+
+    parser.add_argument('--lr', type=float, default = 5e-5)
+    parser.add_argument('--distance_norm_factor', type=float, default = 0.0001)
+    parser.add_argument('--spatial_dist_fill', type=float, default = 20)
+
+    parser.add_argument('--with_type', default=False, action='store_true')
+    parser.add_argument('--sep_between_neighbors', default=False, action='store_true')
+    parser.add_argument('--freeze_backbone', default=False, action='store_true')
+
+    parser.add_argument('--backbone_option', type=str, default='bert-base')
+    parser.add_argument('--model_save_dir', type=str, default=None)
+
+
+                        
+    args = parser.parse_args()
+    print('\n')
+    print(args)
+    print('\n')
+
+
+    # out_dir not None, and out_dir does not exist, then create out_dir
+    if args.model_save_dir is not None and not os.path.isdir(args.model_save_dir):
+        os.makedirs(args.model_save_dir)
+
+    training(args)
+
+
+   
+if __name__ == '__main__':
+
+    main()
+
+    

models/spabert/experiments/semantic_typing/src/train_cls_spatialbert.py

@@ -0,0 +1,276 @@
+import os
+import sys
+from transformers import RobertaTokenizer, BertTokenizer
+from tqdm import tqdm  # for our progress bar
+from transformers import AdamW
+
+import torch
+from torch.utils.data import DataLoader
+
+sys.path.append('../../../')
+from models.spatial_bert_model import SpatialBertModel
+from models.spatial_bert_model import SpatialBertConfig
+from models.spatial_bert_model import  SpatialBertForMaskedLM, SpatialBertForSemanticTyping
+from datasets.osm_sample_loader import PbfMapDataset
+from datasets.const import *
+
+from transformers.models.bert.modeling_bert import BertForMaskedLM
+
+import numpy as np
+import argparse 
+from sklearn.preprocessing import LabelEncoder
+import pdb
+
+
+DEBUG = False
+
+
+def training(args):
+
+    num_workers = args.num_workers
+    batch_size = args.batch_size
+    epochs = args.epochs
+    lr = args.lr #1e-7 # 5e-5
+    save_interval = args.save_interval
+    max_token_len = args.max_token_len
+    distance_norm_factor = args.distance_norm_factor
+    spatial_dist_fill=args.spatial_dist_fill
+    with_type = args.with_type
+    sep_between_neighbors = args.sep_between_neighbors
+    freeze_backbone = args.freeze_backbone
+    mlm_checkpoint_path = args.mlm_checkpoint_path
+
+    if_no_spatial_distance = args.no_spatial_distance
+
+
+    bert_option = args.bert_option
+
+    assert bert_option in ['bert-base','bert-large']
+
+    if args.num_classes == 9:
+        london_file_path = '../../semantic_typing/data/sql_output/osm-point-london-typing.json'
+        california_file_path = '../../semantic_typing/data/sql_output/osm-point-california-typing.json'
+        TYPE_LIST = CLASS_9_LIST
+        type_key_str  = 'class'
+    elif args.num_classes == 74:
+        london_file_path = '../../semantic_typing/data/sql_output/osm-point-london-typing-ranking.json'
+        california_file_path = '../../semantic_typing/data/sql_output/osm-point-california-typing-ranking.json'
+        TYPE_LIST = CLASS_74_LIST
+        type_key_str = 'fine_class'
+    else:
+        raise NotImplementedError
+
+
+    if args.model_save_dir is None:
+        checkpoint_basename = os.path.basename(mlm_checkpoint_path)
+        checkpoint_prefix = checkpoint_basename.replace("mlm_mem_keeppos_","").strip('.pth')
+
+        sep_pathstr = '_sep' if sep_between_neighbors else '_nosep' 
+        freeze_pathstr = '_freeze' if freeze_backbone else '_nofreeze'
+        if if_no_spatial_distance:
+            model_save_dir = '/data2/zekun/spatial_bert_weights_ablation/'
+        else:
+            model_save_dir = '/data2/zekun/spatial_bert_weights/'
+        model_save_dir = os.path.join(model_save_dir, 'typing_lr' + str("{:.0e}".format(lr)) + sep_pathstr +'_'+bert_option+ freeze_pathstr + '_london_california_bsize' + str(batch_size) )
+        model_save_dir = os.path.join(model_save_dir, checkpoint_prefix)
+
+        if not os.path.isdir(model_save_dir):
+            os.makedirs(model_save_dir)
+    else:
+        model_save_dir = args.model_save_dir
+
+    
+    print('model_save_dir', model_save_dir)
+    print('\n')
+
+    if bert_option == 'bert-base':
+        tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
+        config = SpatialBertConfig(use_spatial_distance_embedding = not if_no_spatial_distance, num_semantic_types=len(TYPE_LIST))
+    elif bert_option == 'bert-large':
+        tokenizer = BertTokenizer.from_pretrained("bert-large-uncased")
+        config = SpatialBertConfig(use_spatial_distance_embedding = not if_no_spatial_distance, hidden_size = 1024, intermediate_size = 4096, num_attention_heads=16, num_hidden_layers=24, num_semantic_types=len(TYPE_LIST))
+    else:
+        raise NotImplementedError
+
+
+    
+    label_encoder = LabelEncoder()
+    label_encoder.fit(TYPE_LIST)
+   
+
+    london_train_val_dataset = PbfMapDataset(data_file_path = london_file_path, 
+                                        tokenizer = tokenizer, 
+                                        max_token_len = max_token_len, 
+                                        distance_norm_factor = distance_norm_factor, 
+                                        spatial_dist_fill = spatial_dist_fill, 
+                                        with_type = with_type,
+                                        type_key_str = type_key_str,
+                                        sep_between_neighbors = sep_between_neighbors, 
+                                        label_encoder = label_encoder,
+                                        mode = 'train')
+
+    percent_80 = int(len(london_train_val_dataset) * 0.8)
+    london_train_dataset, london_val_dataset = torch.utils.data.random_split(london_train_val_dataset, [percent_80, len(london_train_val_dataset) - percent_80])
+
+    california_train_val_dataset = PbfMapDataset(data_file_path = california_file_path, 
+                                            tokenizer = tokenizer, 
+                                            max_token_len = max_token_len, 
+                                            distance_norm_factor = distance_norm_factor, 
+                                            spatial_dist_fill = spatial_dist_fill, 
+                                            with_type = with_type,
+                                            type_key_str = type_key_str,
+                                            sep_between_neighbors = sep_between_neighbors,
+                                            label_encoder = label_encoder,
+                                            mode = 'train')
+    percent_80 = int(len(california_train_val_dataset) * 0.8)
+    california_train_dataset, california_val_dataset = torch.utils.data.random_split(california_train_val_dataset, [percent_80, len(california_train_val_dataset) - percent_80])
+
+    train_dataset = torch.utils.data.ConcatDataset([london_train_dataset, california_train_dataset])
+    val_dataset = torch.utils.data.ConcatDataset([london_val_dataset, california_val_dataset])
+
+
+    if DEBUG:
+        train_loader = DataLoader(train_dataset, batch_size= batch_size, num_workers=num_workers,
+                                shuffle=False, pin_memory=True, drop_last=True)
+        val_loader = DataLoader(val_dataset, batch_size= batch_size, num_workers=num_workers,
+                                shuffle=False, pin_memory=True, drop_last=False)
+    else:
+        train_loader = DataLoader(train_dataset, batch_size= batch_size, num_workers=num_workers,
+                                shuffle=True, pin_memory=True, drop_last=True)
+        val_loader = DataLoader(val_dataset, batch_size= batch_size, num_workers=num_workers,
+                                shuffle=False, pin_memory=True, drop_last=False)
+
+
+    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+
+    model = SpatialBertForSemanticTyping(config)
+    model.to(device)
+
+
+    model.load_state_dict(torch.load(mlm_checkpoint_path), strict = False)
+
+    model.train()
+
+    
+
+    # initialize optimizer
+    optim = AdamW(model.parameters(), lr = lr)
+
+    print('start training...')
+
+    for epoch in range(epochs):
+        # setup loop with TQDM and dataloader
+        loop = tqdm(train_loader, leave=True)
+        iter = 0
+        for batch in loop:
+            # initialize calculated gradients (from prev step)
+            optim.zero_grad()
+            # pull all tensor batches required for training
+            input_ids = batch['pseudo_sentence'].to(device)
+            attention_mask = batch['attention_mask'].to(device)
+            position_list_x = batch['norm_lng_list'].to(device)
+            position_list_y = batch['norm_lat_list'].to(device)
+            sent_position_ids = batch['sent_position_ids'].to(device)
+
+            #labels = batch['pseudo_sentence'].to(device)
+            labels = batch['pivot_type'].to(device)
+            pivot_lens = batch['pivot_token_len'].to(device)
+
+            outputs = model(input_ids, attention_mask = attention_mask, sent_position_ids = sent_position_ids,
+                position_list_x = position_list_x, position_list_y = position_list_y, labels = labels, pivot_len_list = pivot_lens)
+
+
+            loss = outputs.loss
+            loss.backward()
+            optim.step()
+
+            loop.set_description(f'Epoch {epoch}')
+            loop.set_postfix({'loss':loss.item()})
+            
+            if DEBUG:
+                print('ep'+str(epoch)+'_' + '_iter'+ str(iter).zfill(5), loss.item() )
+
+            iter += 1
+
+            if iter % save_interval == 0 or iter == loop.total:
+                loss_valid = validating(val_loader, model, device)
+
+                save_path = os.path.join(model_save_dir, 'keeppos_ep'+str(epoch) + '_iter'+ str(iter).zfill(5) \
+                + '_' +str("{:.4f}".format(loss.item())) + '_val' + str("{:.4f}".format(loss_valid)) +'.pth' )
+
+                torch.save(model.state_dict(), save_path)
+                print('validation loss', loss_valid)
+                print('saving model checkpoint to', save_path)
+
+def validating(val_loader, model, device):
+
+    with torch.no_grad():
+
+        loss_valid = 0
+        loop = tqdm(val_loader, leave=True)
+
+        for batch in loop:
+            input_ids = batch['pseudo_sentence'].to(device)
+            attention_mask = batch['attention_mask'].to(device)
+            position_list_x = batch['norm_lng_list'].to(device)
+            position_list_y = batch['norm_lat_list'].to(device)
+            sent_position_ids = batch['sent_position_ids'].to(device)
+
+
+            labels = batch['pivot_type'].to(device)
+            pivot_lens = batch['pivot_token_len'].to(device)
+
+            outputs = model(input_ids, attention_mask = attention_mask, sent_position_ids = sent_position_ids,
+                position_list_x = position_list_x, position_list_y = position_list_y, labels = labels, pivot_len_list = pivot_lens)
+
+            loss_valid  += outputs.loss
+
+        loss_valid /= len(val_loader)
+
+        return loss_valid
+
+
+def main():
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--num_workers', type=int, default=5)
+    parser.add_argument('--batch_size', type=int, default=12)
+    parser.add_argument('--epochs', type=int, default=10)
+    parser.add_argument('--save_interval', type=int, default=2000)
+    parser.add_argument('--max_token_len', type=int, default=512)
+    
+
+    parser.add_argument('--lr', type=float, default = 5e-5)
+    parser.add_argument('--distance_norm_factor', type=float, default = 0.0001)
+    parser.add_argument('--spatial_dist_fill', type=float, default = 100)
+    parser.add_argument('--num_classes', type=int, default = 9)
+
+    parser.add_argument('--with_type', default=False, action='store_true')
+    parser.add_argument('--sep_between_neighbors', default=False, action='store_true')
+    parser.add_argument('--freeze_backbone', default=False, action='store_true')
+    parser.add_argument('--no_spatial_distance',  default=False, action='store_true')
+
+    parser.add_argument('--bert_option', type=str, default='bert-base')
+    parser.add_argument('--model_save_dir', type=str, default=None)
+
+    parser.add_argument('--mlm_checkpoint_path', type=str, default=None)
+
+                        
+    args = parser.parse_args()
+    print('\n')
+    print(args)
+    print('\n')
+
+
+    # out_dir not None, and out_dir does not exist, then create out_dir
+    if args.model_save_dir is not None and not os.path.isdir(args.model_save_dir):
+        os.makedirs(args.model_save_dir)
+
+    training(args)
+
+
+if __name__ == '__main__':
+
+    main()
+
+    

models/spabert/models/baseline_typing_model.py

@@ -0,0 +1,106 @@
+
+import torch
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import CrossEntropyLoss, MSELoss
+from transformers.modeling_outputs import  SequenceClassifierOutput
+
+
+class PivotEntityPooler(nn.Module):
+    def __init__(self):
+        super().__init__()
+        
+
+    def forward(self, hidden_states, pivot_len_list):
+        # We "pool" the model by simply taking the hidden state corresponding
+        # to the tokens of pivot entity
+        
+        bsize = hidden_states.shape[0]
+
+        tensor_list = []
+        for i in torch.arange(0, bsize):
+            pivot_token_full = hidden_states[i, 1:pivot_len_list[i]+1]
+            pivot_token_tensor = torch.mean(torch.unsqueeze(pivot_token_full, 0), dim = 1)
+            tensor_list.append(pivot_token_tensor)
+
+
+        batch_pivot_tensor = torch.cat(tensor_list, dim = 0)
+
+        return batch_pivot_tensor
+       
+
+class BaselineTypingHead(nn.Module):
+    def __init__(self, hidden_size, num_semantic_types):
+        super().__init__()
+        self.dense = nn.Linear(hidden_size, hidden_size)
+        self.activation = nn.Tanh()
+        
+        self.seq_relationship = nn.Linear(hidden_size, num_semantic_types)
+
+    def forward(self, pivot_pooled_output):
+
+        pivot_pooled_output = self.dense(pivot_pooled_output)
+        pivot_pooled_output = self.activation(pivot_pooled_output)
+
+        seq_relationship_score = self.seq_relationship(pivot_pooled_output)
+        return seq_relationship_score
+
+
+class BaselineForSemanticTyping(nn.Module):
+
+    def __init__(self, backbone, hidden_size, num_semantic_types):
+        super().__init__() 
+
+        self.backbone = backbone
+        self.pivot_pooler = PivotEntityPooler() 
+        self.num_semantic_types = num_semantic_types
+
+        self.cls = BaselineTypingHead(hidden_size, num_semantic_types)
+
+
+    def forward(
+        self,
+        input_ids=None,
+        position_ids = None,
+        pivot_len_list = None,
+        attention_mask=None, 
+        head_mask=None,
+        labels=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict = True
+    ):
+
+        outputs = self.backbone(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids = position_ids,
+            head_mask=head_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+        )
+
+
+        sequence_output = outputs[0]
+        pooled_output = self.pivot_pooler(sequence_output, pivot_len_list) 
+
+
+        type_prediction_score = self.cls(pooled_output)
+
+        typing_loss = None
+        if labels is not None:
+            loss_fct = CrossEntropyLoss()
+            typing_loss = loss_fct(type_prediction_score.view(-1, self.num_semantic_types), labels.view(-1))
+
+        
+
+        if not return_dict:
+            output = (type_prediction_score,) + outputs[2:]
+            return ((typing_loss,) + output) if typing_loss is not None else output
+
+        return SequenceClassifierOutput(
+            loss=typing_loss,
+            logits=type_prediction_score,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )