Update README.md

README.md

@@ -31,28 +31,48 @@ model = AutoModelForPreTraining.from_pretrained('felipemaiapolo/legalnlp-bert')
 tokenizer = AutoTokenizer.from_pretrained('felipemaiapolo/legalnlp-bert', do_lower_case=False)
 ```
 
-### Ex. extracting BERT embeddings
+
+### Ex. BERT embeddings
+
+```python
+from transformers import pipeline
+
+pipe = pipeline("feature-extraction", model='felipemaiapolo/legalnlp-bert')
+encoded_sentence = pipe('Juíz negou o recurso.')
+
+
+### Masked language modeling prediction example
 
 ```python
-import torch
-
-model = AutoModel.from_pretrained('felipemaiapolo/legalnlp-bert')
-input_ids = tokenizer.encode('Tinha uma pedra no meio do caminho.', return_tensors='pt')
-
-with torch.no_grad():
-    outs = model(input_ids)
-    encoded = outs[0][0, 1:-1]  # Ignore [CLS] and [SEP] special tokens
-
-# encoded.shape: (8, 768)
-# tensor([[-0.0398, -0.3057,  0.2431,  ..., -0.5420,  0.1857, -0.5775],
-#         [-0.2926, -0.1957,  0.7020,  ..., -0.2843,  0.0530, -0.4304],
-#         [ 0.2463, -0.1467,  0.5496,  ...,  0.3781, -0.2325, -0.5469],
-#         ...,
-#         [ 0.0662,  0.7817,  0.3486,  ..., -0.4131, -0.2852, -0.2819],
-#         [ 0.0662,  0.2845,  0.1871,  ..., -0.2542, -0.2933, -0.0661],
-#         [ 0.2761, -0.1657,  0.3288,  ..., -0.2102,  0.0029, -0.2009]])
+from transformers import pipeline
+
+pipe = pipeline('fill-mask', model=model, tokenizer=tokenizer)
+
+pipe('Juíz negou o [MASK].')
+#  [{'score': 0.6387444734573364,
+#  'token': 7608,
+#  'token_str': 'julgamento',
+#  'sequence': 'juiz negou o julgamento.'},
+# {'score': 0.09632532298564911,
+#  'token': 7509,
+#  'token_str': 'voto',
+#  'sequence': 'juiz negou o voto.'},
+# {'score': 0.06424401700496674,
+#  'token': 17225,
+#  'token_str': 'julgado',
+#  'sequence': 'juiz negou o julgado.'},
+# {'score': 0.05929475650191307,
+#  'token': 8190,
+#  'token_str': 'recurso',
+#  'sequence': 'juiz negou o recurso.'},
+# {'score': 0.011442390270531178,
+#  'token': 6330,
+#  'token_str': 'registro',
+#  'sequence': 'juiz negou o registro.'}]
+
 ```
-    
+
+
 # References
 
 [1] Polo, Felipe Maia, et al. "LegalNLP-Natural Language Processing methods for the Brazilian Legal Language." Anais do XVIII Encontro Nacional de Inteligência Artificial e Computacional. SBC, 2021.