Add gradio app for caption scoring

app.py

@@ -0,0 +1,63 @@
+import sys
+
+import gradio as gr
+import jax
+from huggingface_hub import snapshot_download
+from PIL import Image
+from transformers import AutoTokenizer
+
+LOCAL_PATH = snapshot_download("flax-community/clip-spanish")
+sys.path.append(LOCAL_PATH)
+
+from modeling_hybrid_clip import FlaxHybridCLIP
+from test_on_image import prepare_image, prepare_text
+
+
+def save_file_to_disk(uplaoded_file):
+    temp_file = "/tmp/image.jpeg"
+    im = Image.fromarray(uplaoded_file)
+    im.save(temp_file)
+    # with open(temp_file, "wb") as f:
+    #     f.write(uploaded_file.getbuffer())
+    return temp_file
+
+
+def run_inference(image_path, text, model, tokenizer):
+    pixel_values = prepare_image(image_path, model)
+    input_text = prepare_text(text, tokenizer)
+    model_output = model(
+        input_text["input_ids"],
+        pixel_values,
+        attention_mask=input_text["attention_mask"],
+        train=False,
+        return_dict=True,
+    )
+    logits = model_output["logits_per_image"]
+    score = jax.nn.sigmoid(logits)[0][0]
+    return score
+
+
+def load_tokenizer_and_model():
+    # load the saved model
+    tokenizer = AutoTokenizer.from_pretrained(
+        "bertin-project/bertin-roberta-base-spanish"
+    )
+    model = FlaxHybridCLIP.from_pretrained(LOCAL_PATH)
+    return tokenizer, model
+
+
+tokenizer, model = load_tokenizer_and_model()
+
+
+def score_image_caption_pair(uploaded_file, text_input):
+    local_image_path = save_file_to_disk(uploaded_file)
+    score = run_inference(
+        local_image_path, text_input, model, tokenizer).tolist()
+    return {"Score": score}
+
+
+image = gr.inputs.Image(shape=(299, 299))
+iface = gr.Interface(
+    fn=score_image_caption_pair, inputs=[image, "text"], outputs="label"
+)
+iface.launch()

requirements.txt

@@ -0,0 +1,9 @@
+flax==0.3.4
+gradio==2.2.2
+huggingface-hub==0.0.12
+jax==0.2.17
+streamlit==0.84.1
+torch==1.9.0
+torchvision==0.10.0
+transformers==4.8.2
+watchdog==2.1.3