diff --git a/lib/ngram_index.py b/lib/ngram_index.py index b4617187cd1c73f355711b9f25660455be904a74..8bf4ce3c87e5b3d286ad03810ed168742310fb07 100644 --- a/lib/ngram_index.py +++ b/lib/ngram_index.py @@ -78,13 +78,23 @@ class NgramIndex(): else: self.freq_ngram[ngram] += 1 - def filter_ngram(self,threshold=20): + def filter_ngram_by_freq(self,threshold=20): freq_data = pd.DataFrame(self.freq_ngram.items(),columns="ngram freq".split()) selected_ngram = freq_data[freq_data.freq<threshold].ngram.values for ng in selected_ngram: index = self.ngram_index[ng] del self.ngram_index[ng] del self.index_ngram[index] + + def filter_top_ngram(self,threshold=20000): + freq_data = pd.DataFrame(self.freq_ngram.items(),columns="ngram freq".split()).sort_values(by="freq",ascending=False) + if len(self.ngram_index)-threshold <0: + return 0 + selected_ngram = freq_data.tail(len(self.ngram_index)-threshold).ngram.values + for ng in selected_ngram: + index = self.ngram_index[ng] + del self.ngram_index[ng] + del self.index_ngram[index] def encode(self,word,complete=True): diff --git a/train_geocoder.py b/train_geocoder.py index d143b58e65f054b8274d51c55b66bb8000158c90..b8c6978c6f2d8991c911bdb6516fdcb619db2148 100644 --- a/train_geocoder.py +++ b/train_geocoder.py @@ -112,7 +112,7 @@ if args.tokenization_method == "bert": # Identify all ngram available pairs_of_toponym.toponym.apply(lambda x : index.split_and_add(x)) pairs_of_toponym.toponym_context.apply(lambda x : index.split_and_add(x)) -index.filter_ngram() +index.filter_top_ngram(40000) num_words = len(index.index_ngram) # necessary for the embedding matrix