Update README.md

README.md

@@ -14,31 +14,34 @@ The Document Image Transformer (DiT) is a transformer encoder model (BERT-like)
 
 Images are presented to the model as a sequence of fixed-size patches (resolution 16x16), which are linearly embedded. One also adds absolute position embeddings before feeding the sequence to the layers of the Transformer encoder.
 
-Note that this model does not provide any fine-tuned heads, hence it's meant to be fine-tuned on tasks like document image classification, table detection or document layout analysis. 
-
 By pre-training the model, it learns an inner representation of images that can then be used to extract features useful for downstream tasks: if you have a dataset of labeled document images for instance, you can train a standard classifier by placing a linear layer on top of the pre-trained encoder.
 
 ## Intended uses & limitations
 
-You can use the raw model for encoding document images into a vector space, but it's mostly meant to be fine-tuned. See the [model hub](https://huggingface.co/models?search=microsoft/dit) to look for fine-tuned versions on a task that interests you.
+You can use the raw model for encoding document images into a vector space, but it's mostly meant to be fine-tuned on tasks like document image classification, table detection or document layout analysis. See the [model hub](https://huggingface.co/models?search=microsoft/dit) to look for fine-tuned versions on a task that interests you.
 
 ### How to use
 
 Here is how to use this model in PyTorch:
 
 ```python
-from transformers import AutoFeatureExtractor, AutoModel
+from transformers import BeitFeatureExtractor, BeitForMaskedImageModeling
+import torch
 from PIL import Image
 
 image = Image.open('path_to_your_document_image')
 
-feature_extractor = AutoFeatureExtractor.from_pretrained('microsoft/dit-base')
-model = AutoModel.from_pretrained('microsoft/dit-base')
+feature_extractor = BeitFeatureExtractor.from_pretrained("microsoft/dit-base")
+model = BeitForMaskedImageModeling.from_pretrained("microsoft/dit-base")
 
-inputs = feature_extractor(images=image, return_tensors="pt")
+num_patches = (model.config.image_size // model.config.patch_size) ** 2
+pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values
+# create random boolean mask of shape (batch_size, num_patches)
+bool_masked_pos = torch.randint(low=0, high=2, size=(1, num_patches)).bool()
 
-outputs = model(**inputs)
-last_hidden_states = outputs.last_hidden_state
+outputs = model(pixel_values, bool_masked_pos=bool_masked_pos)
+loss, logits = outputs.loss, outputs.logits
+list(logits.shape)
 ```
 
 ### BibTeX entry and citation info