Update README.md

README.md

@@ -51,27 +51,32 @@ In order to use SEC-BERT-NUM, you have to pre-process texts replacing every nume
 Below there is an example how you can pre-process a simple sentence. This approach is quite simple, feel free to modify it as you see fit.
 
 ```python
+import re
 import spacy
 from transformers import AutoTokenizer
 
-tokenizer = AutoTokenizer.from_pretrained("nlpaueb/sec-bert-num")
+tokenizer = AutoTokenizer.from_pretrained("nlpaueb/sec-bert-shape")
 spacy_tokenizer = spacy.load("en_core_web_sm")
 
 sentence = "Total net sales decreased 2% or $5.4 billion during 2019 compared to 2018."
-tokens = [t.text for t in spacy_tokenizer(sentence)]
-
-processed_sentence = []
-for token in tokens:
-    if re.fullmatch(r"(\d+[\d,.]*)|([,.]\d+)", token):
-        shape = '[' + re.sub(r'\d', 'X', token) + ']'
-        if shape in tokenizer.additional_special_tokens:
-            processed_sentence.append(shape)
+
+def sec_bert_shape_preprocess(text):
+    tokens = [t.text for t in spacy_tokenizer(sentence)]
+
+    processed_text = []
+    for token in tokens:
+        if re.fullmatch(r"(\d+[\d,.]*)|([,.]\d+)", token):
+            shape = '[' + re.sub(r'\d', 'X', token) + ']'
+            if shape in tokenizer.additional_special_tokens:
+                processed_text.append(shape)
+            else:
+                processed_text.append('[NUM]')
         else:
-            processed_sentence.append('[NUM]')
-    else:
-        processed_sentence.append(token)
+            processed_text.append(token)
+            
+    return ' '.join(processed_text)
         
-tokenized_sentence = tokenizer.tokenize(' '.join(processed_sentence))
+tokenized_sentence = tokenizer.tokenize(sec_bert_shape_preprocess(sentence))
 print(tokenized_sentence)
 """
 ['total', 'net', 'sales', 'decreased', '[X]', '%', 'or', '$', '[X.X]', 'billion', 'during', '[XXXX]', 'compared', 'to', '[XXXX]', '.']