Add MLM task

app.py

@@ -1,183 +1,52 @@
-import json
+from apps import mlm, vqa
 import os
-from io import BytesIO
-
-import matplotlib.pyplot as plt
-import numpy as np
-import pandas as pd
 import streamlit as st
-from mtranslate import translate
-from PIL import Image
-from streamlit.elements import markdown
-
-from model.flax_clip_vision_bert.modeling_clip_vision_bert import (
-    FlaxCLIPVisionBertForSequenceClassification,
-)
-from session import _get_state
-from utils import (
-    get_text_attributes,
-    get_top_5_predictions,
-    get_transformed_image,
-    plotly_express_horizontal_bar_plot,
-    translate_labels,
-)
-
-state = _get_state()
-
-
-@st.cache(persist=True)
-def load_model(ckpt):
-    return FlaxCLIPVisionBertForSequenceClassification.from_pretrained(ckpt)
-
-
-@st.cache(persist=True)
-def predict(transformed_image, question_inputs):
-    return np.array(model(pixel_values=transformed_image, **question_inputs)[0][0])
-
-
-def softmax(logits):
-    return np.exp(logits) / np.sum(np.exp(logits), axis=0)
-
+from multiapp import MultiApp
 
 def read_markdown(path, parent="./sections/"):
     with open(os.path.join(parent, path)) as f:
         return f.read()
 
-checkpoints = ["./ckpt/vqa/ckpt-60k-5999"]  # TODO: Maybe add more checkpoints?
-dummy_data = pd.read_csv("dummy_vqa_multilingual.tsv", sep="\t")
-code_to_name = {
-    "en": "English",
-    "fr": "French",
-    "de": "German",
-    "es": "Spanish",
-}
-
-with open("answer_reverse_mapping.json") as f:
-    answer_reverse_mapping = json.load(f)
-
-
-st.set_page_config(
-    page_title="Multilingual VQA",
-    layout="wide",
-    initial_sidebar_state="collapsed",
-    page_icon="./misc/mvqa-logo-3-white.png",
-)
-
-st.title("Multilingual Visual Question Answering")
-st.write(
-    "[Gunjan Chhablani](https://huggingface.co/gchhablani), [Bhavitvya Malik](https://huggingface.co/bhavitvyamalik)"
-)
-
-image_col, intro_col = st.beta_columns([3, 8])
-image_col.image("./misc/mvqa-logo-3-white.png", use_column_width="always")
-intro_col.write(read_markdown("intro.md"))
-with st.beta_expander("Usage"):
-    st.write(read_markdown("usage.md"))
-
-with st.beta_expander("Article"):
-    st.write(read_markdown("abstract.md"))
-    st.write(read_markdown("caveats.md"))
-    st.write("## Methodology")
-    col1, col2 = st.beta_columns([1,1])
-    col1.image(
-        "./misc/article/resized/Multilingual-VQA.png",
-        caption="Masked LM model for Image-text Pretraining.",
+def main():
+    st.set_page_config(
+        page_title="Multilingual VQA",
+        layout="wide",
+        initial_sidebar_state="collapsed",
+        page_icon="./misc/mvqa-logo-3-white.png",
     )
-    col2.markdown(read_markdown("pretraining.md"))
-    st.markdown(read_markdown("finetuning.md"))
-    st.write(read_markdown("challenges.md"))
-    st.write(read_markdown("social_impact.md"))
-    st.write(read_markdown("references.md"))
-    st.write(read_markdown("checkpoints.md"))
-    st.write(read_markdown("acknowledgements.md"))
 
-first_index = 20
-# Init Session State
-if state.image_file is None:
-    state.image_file = dummy_data.loc[first_index, "image_file"]
-    state.question = dummy_data.loc[first_index, "question"].strip("- ")
-    state.answer_label = dummy_data.loc[first_index, "answer_label"]
-    state.question_lang_id = dummy_data.loc[first_index, "lang_id"]
-    state.answer_lang_id = dummy_data.loc[first_index, "lang_id"]
-
-    image_path = os.path.join("resized_images", state.image_file)
-    image = plt.imread(image_path)
-    state.image = image
-
-# col1, col2, col3 = st.beta_columns([3,3,3])
-
-if st.button(
-    "Get a random example",
-    help="Get a random example from the 100 `seeded` image-text pairs.",
-):
-    sample = dummy_data.sample(1).reset_index()
-    state.image_file = sample.loc[0, "image_file"]
-    state.question = sample.loc[0, "question"].strip("- ")
-    state.answer_label = sample.loc[0, "answer_label"]
-    state.question_lang_id = sample.loc[0, "lang_id"]
-    state.answer_lang_id = sample.loc[0, "lang_id"]
-
-    image_path = os.path.join("resized_images", state.image_file)
-    image = plt.imread(image_path)
-    state.image = image
-
-# col2.write("OR")
-
-# uploaded_file = col2.file_uploader(
-#     "Upload your image",
-#     type=["png", "jpg", "jpeg"],
-#     help="Upload a file of your choosing.",
-# )
-# if uploaded_file is not None:
-#     state.image_file = os.path.join("images/val2014", uploaded_file.name)
-#     state.image = np.array(Image.open(uploaded_file))
-
-transformed_image = get_transformed_image(state.image)
-
-new_col1, new_col2 = st.beta_columns([5, 5])
-
-# Display Image
-new_col1.image(state.image, use_column_width="always")
-
-
-# Display Question
-question = new_col2.text_input(
-    label="Question",
-    value=state.question,
-    help="Type your question regarding the image above in one of the four languages.",
-)
-new_col2.markdown(
-    f"""**English Translation**: {question if state.question_lang_id == "en" else translate(question, 'en')}"""
-)
-
-question_inputs = get_text_attributes(question)
-
-# Select Language
-options = ["en", "de", "es", "fr"]
-state.answer_lang_id = new_col2.selectbox(
-    "Answer Language",
-    index=options.index(state.answer_lang_id),
-    options=options,
-    format_func=lambda x: code_to_name[x],
-    help="The language to be used to show the top-5 labels.",
-)
-
-actual_answer = answer_reverse_mapping[str(state.answer_label)]
-new_col2.markdown(
-    "**Actual Answer**: "
-    + translate_labels([actual_answer], state.answer_lang_id)[0]
-    + " ("
-    + actual_answer
-    + ")"
-)
+    st.title("Multilingual Visual Question Answering")
+    st.write(
+        "[Gunjan Chhablani](https://huggingface.co/gchhablani), [Bhavitvya Malik](https://huggingface.co/bhavitvyamalik)"
+    )
 
-# Display Top-5 Predictions
-with st.spinner("Loading model..."):
-    model = load_model(checkpoints[0])
-with st.spinner("Predicting..."):
-    logits = predict(transformed_image, dict(question_inputs))
-logits = softmax(logits)
-labels, values = get_top_5_predictions(logits, answer_reverse_mapping)
-translated_labels = translate_labels(labels, state.answer_lang_id)
-fig = plotly_express_horizontal_bar_plot(values, translated_labels)
-st.plotly_chart(fig, use_container_width=True)
+    image_col, intro_col = st.beta_columns([3, 8])
+    image_col.image("./misc/mvqa-logo-3-white.png", use_column_width="always")
+    intro_col.write(read_markdown("intro.md"))
+    with st.beta_expander("Usage"):
+        st.write(read_markdown("usage.md"))
+
+    with st.beta_expander("Article"):
+        st.write(read_markdown("abstract.md"))
+        st.write(read_markdown("caveats.md"))
+        st.write("## Methodology")
+        col1, col2 = st.beta_columns([1,1])
+        col1.image(
+            "./misc/article/Multilingual-VQA.png",
+            caption="Masked LM model for Image-text Pretraining.",
+        )
+        col2.markdown(read_markdown("pretraining.md"))
+        st.markdown(read_markdown("finetuning.md"))
+        st.write(read_markdown("challenges.md"))
+        st.write(read_markdown("social_impact.md"))
+        st.write(read_markdown("references.md"))
+        st.write(read_markdown("checkpoints.md"))
+        st.write(read_markdown("acknowledgements.md"))
+
+    app = MultiApp()
+    app.add_app("Visual Question Answering", vqa.app)
+    app.add_app("Mask Filling", mlm.app)
+    app.run()
+
+if __name__ == "__main__":
+    main()

apps/mlm.py

@@ -0,0 +1,109 @@
+
+from .utils import (
+    get_text_attributes,
+    get_top_5_predictions,
+    get_transformed_image,
+    plotly_express_horizontal_bar_plot,
+    translate_labels,
+    bert_tokenizer
+)
+
+import streamlit as st
+import numpy as np
+import pandas as pd
+import os
+import matplotlib.pyplot as plt
+
+from session import _get_state
+
+
+from .model.flax_clip_vision_bert.modeling_clip_vision_bert import (
+    FlaxCLIPVisionBertForMaskedLM,
+)
+
+def softmax(logits):
+    return np.exp(logits) / np.sum(np.exp(logits), axis=0)
+
+def app():
+    state = _get_state()
+
+    @st.cache(persist=False)
+    def predict(transformed_image, caption_inputs):
+        outputs = state.model(pixel_values=transformed_image, **caption_inputs)
+        indices = np.where(caption_inputs['input_ids']==bert_tokenizer.mask_token_id)
+        preds = outputs.logits[indices][0]
+        sorted_indices = np.argsort(preds)[::-1] # Get reverse sorted scores
+        top_5_indices = sorted_indices[:5]
+        top_5_tokens = bert_tokenizer.convert_ids_to_tokens(top_5_indices)
+        top_5_scores = np.array(preds[top_5_indices])
+        return top_5_tokens, top_5_scores
+
+
+    @st.cache(persist=False)
+    def load_model(ckpt):
+        return FlaxCLIPVisionBertForMaskedLM.from_pretrained(ckpt)
+
+    mlm_checkpoints = ['flax-community/clip-vision-bert-cc12m-70k']
+    dummy_data = pd.read_csv("cc12m_data/vqa_val.tsv", sep="\t")
+
+    first_index = 20
+    # Init Session State
+    if state.image_file is None:
+        state.image_file = dummy_data.loc[first_index, "image_file"]
+        caption = dummy_data.loc[first_index, "caption"].strip("- ")
+        ids = bert_tokenizer(caption)
+        ids[np.random.randint(0, len(ids))] = bert_tokenizer.mask_token_id
+        state.caption = bert_tokenizer.decode(ids)
+        state.caption_lang_id = dummy_data.loc[first_index, "lang_id"]
+
+        image_path = os.path.join("cc12m_data/images_vqa", state.image_file)
+        image = plt.imread(image_path)
+        state.image = image
+
+    if state.model is None:
+        # Display Top-5 Predictions
+        with st.spinner("Loading model..."):
+            state.model = load_model(mlm_checkpoints[0])
+
+    if st.button(
+        "Get a random example",
+        help="Get a random example from the 100 `seeded` image-text pairs.",
+    ):
+        sample = dummy_data.sample(1).reset_index()
+        state.image_file = sample.loc[0, "image_file"]
+        caption = sample.loc[0, "caption"].strip("- ")
+        ids = bert_tokenizer(caption)
+        ids[np.random.randint(0, len(ids))] = bert_tokenizer.mask_token_id
+        state.caption = bert_tokenizer.decode(ids)
+        state.caption_lang_id = sample.loc[0, "lang_id"]
+
+        image_path = os.path.join("cc12m_data/images_vqa", state.image_file)
+        image = plt.imread(image_path)
+        state.image = image
+
+    transformed_image = get_transformed_image(state.image)
+
+    new_col1, new_col2 = st.beta_columns([5, 5])
+
+    # Display Image
+    new_col1.image(state.image, use_column_width="always")
+
+
+    # Display caption
+    new_col2.write("Write your text with exactly one [MASK] token.")
+    caption = new_col2.text_input(
+        label="Text",
+        value=state.caption,
+        help="Type your masked caption regarding the image above in one of the four languages.",
+    )
+
+    caption_inputs = get_text_attributes(caption)
+
+    # Display Top-5 Predictions
+
+    with st.spinner("Predicting..."):
+        logits = predict(transformed_image, dict(caption_inputs))
+    logits = softmax(logits)
+    labels, values = get_top_5_predictions(logits)
+    fig = plotly_express_horizontal_bar_plot(values, labels)
+    st.plotly_chart(fig, use_container_width=True)

utils.py → apps/utils.py

@@ -3,8 +3,7 @@ import json
 import numpy as np
 import plotly.express as px
 import torch
-from PIL import Image
-from torchvision.io import ImageReadMode, read_image
+from torchvision.io import read_image
 from torchvision.transforms import CenterCrop, ConvertImageDtype, Normalize, Resize
 from torchvision.transforms.functional import InterpolationMode
 from transformers import BertTokenizerFast
@@ -41,15 +40,17 @@ def get_transformed_image(image):
 
 bert_tokenizer = BertTokenizerFast.from_pretrained("bert-base-multilingual-uncased")
 
-
 def get_text_attributes(text):
     return bert_tokenizer([text], return_token_type_ids=True, return_tensors="np")
 
 
-def get_top_5_predictions(logits, answer_reverse_mapping):
+def get_top_5_predictions(logits, answer_reverse_mapping=None):
     indices = np.argsort(logits)[-5:]
     values = logits[indices]
-    labels = [answer_reverse_mapping[str(i)] for i in indices]
+    if answer_reverse_mapping is not None:
+        labels = [answer_reverse_mapping[str(i)] for i in indices]
+    else:
+        labels = bert_tokenizer.convert_ids_to_tokens(indices)
     return labels, values
 
 

apps/vqa.py

@@ -0,0 +1,131 @@
+
+from .utils import (
+    get_text_attributes,
+    get_top_5_predictions,
+    get_transformed_image,
+    plotly_express_horizontal_bar_plot,
+    translate_labels,
+)
+
+import streamlit as st
+import numpy as np
+import pandas as pd
+import os
+import matplotlib.pyplot as plt
+import json
+
+from mtranslate import translate
+from session import _get_state
+
+
+from .model.flax_clip_vision_bert.modeling_clip_vision_bert import (
+    FlaxCLIPVisionBertForSequenceClassification,
+)
+
+def softmax(logits):
+    return np.exp(logits) / np.sum(np.exp(logits), axis=0)
+
+def app():
+    state = _get_state()
+
+    @st.cache(persist=True)
+    def predict(transformed_image, question_inputs):
+        return np.array(state.model(pixel_values=transformed_image, **question_inputs)[0][0])
+
+
+    @st.cache(persist=True)
+    def load_model(ckpt):
+        return FlaxCLIPVisionBertForSequenceClassification.from_pretrained(ckpt)
+
+    vqa_checkpoints = ["flax-community/clip-vision-bert-vqa-ft-6k"]  # TODO: Maybe add more checkpoints?
+    dummy_data = pd.read_csv("dummy_vqa_multilingual.tsv", sep="\t")
+    code_to_name = {
+        "en": "English",
+        "fr": "French",
+        "de": "German",
+        "es": "Spanish",
+    }
+
+
+    with open("answer_reverse_mapping.json") as f:
+        answer_reverse_mapping = json.load(f)
+
+    first_index = 20
+    # Init Session State
+    if state.image_file is None:
+        state.image_file = dummy_data.loc[first_index, "image_file"]
+        state.question = dummy_data.loc[first_index, "question"].strip("- ")
+        state.answer_label = dummy_data.loc[first_index, "answer_label"]
+        state.question_lang_id = dummy_data.loc[first_index, "lang_id"]
+        state.answer_lang_id = dummy_data.loc[first_index, "lang_id"]
+
+        image_path = os.path.join("resized_images", state.image_file)
+        image = plt.imread(image_path)
+        state.image = image
+
+    if state.model is None:
+        # Display Top-5 Predictions
+        with st.spinner("Loading model..."):
+            state.model = load_model(vqa_checkpoints[0])
+
+    if st.button(
+        "Get a random example",
+        help="Get a random example from the 100 `seeded` image-text pairs.",
+    ):
+        sample = dummy_data.sample(1).reset_index()
+        state.image_file = sample.loc[0, "image_file"]
+        state.question = sample.loc[0, "question"].strip("- ")
+        state.answer_label = sample.loc[0, "answer_label"]
+        state.question_lang_id = sample.loc[0, "lang_id"]
+        state.answer_lang_id = sample.loc[0, "lang_id"]
+
+        image_path = os.path.join("resized_images", state.image_file)
+        image = plt.imread(image_path)
+        state.image = image
+
+    transformed_image = get_transformed_image(state.image)
+
+    new_col1, new_col2 = st.beta_columns([5, 5])
+
+    # Display Image
+    new_col1.image(state.image, use_column_width="always")
+
+
+    # Display Question
+    question = new_col2.text_input(
+        label="Question",
+        value=state.question,
+        help="Type your question regarding the image above in one of the four languages.",
+    )
+    new_col2.markdown(
+        f"""**English Translation**: {question if state.question_lang_id == "en" else translate(question, 'en')}"""
+    )
+
+    question_inputs = get_text_attributes(question)
+
+    # Select Language
+    options = ["en", "de", "es", "fr"]
+    state.answer_lang_id = new_col2.selectbox(
+        "Answer Language",
+        index=options.index(state.answer_lang_id),
+        options=options,
+        format_func=lambda x: code_to_name[x],
+        help="The language to be used to show the top-5 labels.",
+    )
+
+    actual_answer = answer_reverse_mapping[str(state.answer_label)]
+    new_col2.markdown(
+        "**Actual Answer**: "
+        + translate_labels([actual_answer], state.answer_lang_id)[0]
+        + " ("
+        + actual_answer
+        + ")"
+    )
+
+    with st.spinner("Predicting..."):
+        logits = predict(transformed_image, dict(question_inputs))
+    logits = softmax(logits)
+    labels, values = get_top_5_predictions(logits, answer_reverse_mapping)
+    translated_labels = translate_labels(labels, state.answer_lang_id)
+    fig = plotly_express_horizontal_bar_plot(values, translated_labels)
+    st.plotly_chart(fig, use_container_width=True)