Update examples in README.md

README.md

@@ -28,42 +28,54 @@ You can use the raw model for image and text retrieval.
 
 ### How to use
 
-Here is how to use this model to get the features of a given text in PyTorch:
+Here is how to use this model to perform image and text matching:
+
 ```python
-import os
-from PIL import Image
-from glob import glob
-from tqdm import tqdm
-import torch
 from transformers import BridgeTowerProcessor, BridgeTowerForImageAndTextRetrieval
+import requests
+from PIL import Image
 
-image_dir = "/datasets/COCO2017/val2017"
-search_text = "a woman holding an umbrella"
+url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+image = Image.open(requests.get(url, stream=True).raw)
+texts = ["An image of two cats chilling on a couch", "A football player scoring a goal"]
 
-processor = BridgeTowerProcessor.from_pretrained(("BridgeTower/bridgetower-base-itm-mlm"))
+processor = BridgeTowerProcessor.from_pretrained("BridgeTower/bridgetower-base-itm-mlm")
 model = BridgeTowerForImageAndTextRetrieval.from_pretrained("BridgeTower/bridgetower-base-itm-mlm")
 
-max_score = float('-inf')
-best_match_image = None
-image_paths = glob(os.path.join(image_dir, '*.jpg'))[:1000]
+# forward pass
+scores = dict()
+for text in texts:
+    # prepare inputs
+    encoding = processor(image, text, return_tensors="pt")
+    outputs = model(**encoding)
+    scores[text] = outputs.logits[0, :].item()
+```
 
-for image_path in tqdm(image_paths, smoothing=1):
-    image = Image.open(image_path).convert("RGB")
-    inputs = processor(image, search_text, return_tensors="pt")
+Here is how to use this model to perfom masked language modeling:
 
-    inputs = dict((k,v.to(device)) if isinstance(v, torch.Tensor) else (k,v) for k,v in inputs.items())
+```python
+from transformers import BridgeTowerProcessor, BridgeTowerForMaskedLM
+from PIL import Image
+
+url = "http://images.cocodataset.org/val2017/000000360943.jpg"
+image = Image.open(requests.get(url, stream=True).raw).convert("RGB")
+text = "a <mask> looking out of the window"
 
-    outputs = model(**inputs)
+processor = BridgeTowerProcessor.from_pretrained(("BridgeTower/bridgetower-base-itm-mlm"))
+model = BridgeTowerForMaskedLM.from_pretrained("BridgeTower/bridgetower-base-itm-mlm")
 
-    score = outputs.logits[0,1].item()
+# prepare inputs
+encoding = processor(image, text, return_tensors="pt")
 
-    if score > max_score:
-        max_score = score 
-        best_match_image = image_path
+# forward pass
+outputs = model(**encoding)
 
-print(max_score)
-print(best_match_image)
+results = processor.decode(outputs.logits.argmax(dim=-1).squeeze(0).tolist())
+
+print(results)
+a cat looking out of the window.
 ```
+
 ### Limitations and bias
 
 TODO