from io import BytesIO import streamlit as st import pandas as pd import json import os import numpy as np from model.flax_clip_vision_bert.modeling_clip_vision_bert import FlaxCLIPVisionBertForSequenceClassification from utils import get_transformed_image, get_text_attributes, get_top_5_predictions, plotly_express_horizontal_bar_plot, translate_labels import matplotlib.pyplot as plt from mtranslate import translate from PIL import Image from session import _get_state state = _get_state() @st.cache def load_model(ckpt): return FlaxCLIPVisionBertForSequenceClassification.from_pretrained(ckpt) def softmax(logits): return np.exp(logits)/np.sum(np.exp(logits), axis=0) checkpoints = ['./ckpt/ckpt-60k-5999'] # TODO: Maybe add more checkpoints? dummy_data = pd.read_csv('dummy_vqa_multilingual.tsv', sep='\t') 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.png", ) st.title("Multilingual Visual Question Answering") with st.beta_expander("About"): st.write("This project is focused on Mutilingual Visual Question Answering. Most of the existing datasets and models on this task work with English-only image-text pairs. Our intention here is to provide a Proof-of-Concept with our simple ViT+BERT model which can be trained on multilingual text checkpoints with pre-trained image encoders and well enough. Due to lack of good-quality multilingual data, we translate subsets of the Conceptual 12M dataset into English (already in English), French, German and Spanish using the Marian models. We achieved 0.49 accuracy on the multilingual validation set we created. With better captions, and hyperparameter-tuning, we expect to see higher performance.") with st.beta_expander("Method"): col1, col2 = st.beta_columns([5,4]) col1.image("./misc/Multilingual-VQA.png") col2.markdown(""" ## Pretraining We follow an approach similar to [VisualBERT](https://arxiv.org/abs/1908.03557). Instead of using a FasterRCNN to get image features, we use a ViT encoder. The task is text-only MLM (Masked Language Modeling). We mask only the text tokens and try to predict the masked tokens. The VisualBERT authors also use a sentence-image matching task where two captions are matched against an image, but we skip this for the sake of simplicity. ### Dataset The dataset we use for pre-training is a cleaned version of [Conceptual 12M](https://github.com/google-research-datasets/conceptual-12m). The dataset is downloaded and then broken images are removed which gives us about 10M images. Then we use the MBart50 ` mbart-large-50-one-to-many-mmt` checkpoint to translate the dataset into four different languages - English, French, German, and Spanish, keeping 2.5 million examples of each language. """) st.markdown(""" ### Model The model is shown in the image above.We create a custom model in Flax which integerates the ViT model inside BERT embeddings. We also use custom configs and modules in order to accomodate for these changes, and allow loading from BERT and ViT checkpoints. The image is fed to the ViT encoder and the text is fed to the word-embedding layers of BERT model. We use the `bert-base-multilingual-uncased` and `openai/clip-vit-base-patch32` checkpoints for BERT and ViT (actually CLIPVision) models, respectively. All our code is available on [GitHub](https://github.com/gchhablani/multilingual-vqa). ## Fine-tuning ### Dataset For fine-tuning, we use the [VQA 2.0](https://visualqa.org/) dataset - particularly, the `train` and `validation` sets. We translate all the questions into the four languages specified above using language-specific MarianMT models. This is because MarianMT models return better labels and are faster, hence, are better for fine-tuning. We get 4x the number of examples in each subset. ### Model We use the `SequenceClassification` model as reference to create our own sequence classification model. 3129 answer labels are chosen, as is the convention for the English VQA task, which can be found [here](https://github.com/gchhablani/multilingual-vqa/blob/main/answer_mapping.json). These are the same labels used in fine-tuning of the VisualBERT models. The outputs shown here have been translated using the [`mtranslate`](https://github.com/mouuff/mtranslate) Google Translate API library. Then we use various pre-trained checkpoints and train the sequence classification model for various steps. Checkpoints: - Pre-trained checkpoint: [multilingual-vqa](https://huggingface.co/flax-community/multilingual-vqa) - Fine-tuned on 45k pretrained checkpoint: [multilingual-vqa-pt-45k-ft](https://huggingface.co/flax-community/multilingual-vqa-pt-45k-ft) - Fine-tuned on 45k pretrained checkpoint with AdaFactor (others use AdamW): [multilingual-vqa-pt-45k-ft-adf](https://huggingface.co/flax-community/multilingual-vqa-pt-45k-ft-adf) - Fine-tuned on 60k pretrained checkpoint: [multilingual-vqa-pt-60k-ft](https://huggingface.co/flax-community/multilingual-vqa-pt-60k-ft) - Fine-tuned on 70k pretrained checkpoint: [multilingual-vqa-pt-60k-ft](https://huggingface.co/flax-community/multilingual-vqa-pt-70k-ft) - From scratch (without pre-training) model: [multilingual-vqa-ft](https://huggingface.co/flax-community/multilingual-vqa-ft) **Caveat**: The best fine-tuned model only achieves 0.49 accuracy on the multilingual validation data that we create. This could be because of not-so-great quality translations, sub-optimal hyperparameters and lack of ample training. In future, we hope to improve this model by addressing such concerns. """) with st.beta_expander("Cherry-Picked Results"): pass with st.beta_expander("Conclusion"): pass with st.beta_expander("Usage"): pass # Init Session State if state.image_file is None: state.image_file = dummy_data.loc[0,'image_file'] state.question = dummy_data.loc[0,'question'].strip('- ') state.answer_label = dummy_data.loc[0,'answer_label'] state.question_lang_id = dummy_data.loc[0, 'lang_id'] state.answer_lang_id = dummy_data.loc[0, 'lang_id'] image_path = os.path.join('images',state.image_file) image = plt.imread(image_path) state.image = image col1, col2 = st.beta_columns([5,5]) # Display Image col1.image(state.image, use_column_width='always') if col2.button('Get a random example'): 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('images',state.image_file) image = plt.imread(image_path) state.image = image st.write("OR") uploaded_file = col2.file_uploader('Upload your image', type=['png','jpg','jpeg']) 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) # Display Question question = st.text_input(label="Question", value=state.question) st.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 = st.selectbox('Answer Language', index=options.index(state.answer_lang_id), options=options) # Display Top-5 Predictions with st.spinner('Loading model...'): model = load_model(checkpoints[0]) with st.spinner('Predicting...'): predictions = model(pixel_values = transformed_image, **question_inputs) logits = np.array(predictions[0][0]) 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)