add clip-base app

app.py

@@ -0,0 +1,81 @@
+import os
+import sys
+import jax
+import streamlit as st
+import transformers
+from huggingface_hub import snapshot_download
+from transformers import AutoTokenizer
+import torch
+from torchvision.io import ImageReadMode, read_image
+
+
+LOCAL_PATH = snapshot_download("flax-community/medclip")
+sys.path.append(LOCAL_PATH)
+
+from src.modeling_medclip import FlaxMedCLIP
+
+def prepare_image(image_path, model):
+    image = read_image(image_path, mode=ImageReadMode.RGB)
+    preprocess = Transform(model.config.vision_config.image_size)
+    preprocess = torch.jit.script(preprocess)
+    preprocessed_image = preprocess(image)
+    pixel_values = torch.stack([preprocessed_image]).permute(0, 2, 3, 1).numpy()
+    return pixel_values
+
+def prepare_text(text, tokenizer):
+    return tokenizer(text, return_tensors="np")
+
+def save_file_to_disk(uplaoded_file):
+    temp_file = os.path.join("/tmp", uplaoded_file.name)
+    with open(temp_file, "wb") as f:
+        f.write(uploaded_file.getbuffer())
+    return temp_file
+@st.cache(
+    hash_funcs={
+        transformers.models.roberta.tokenization_roberta_fast.RobertaTokenizerFast: id,
+        FlaxHybridCLIP: id,
+    },
+    show_spinner=False
+)
+def load_tokenizer_and_model():
+    # load the saved model
+    tokenizer = AutoTokenizer.from_pretrained("allenai/scibert_scivocab_uncased")
+    model = FlaxHybridCLIP.from_pretrained(LOCAL_PATH)
+    return tokenizer, model
+
+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
+
+tokenizer, model = load_tokenizer_and_model()
+st.title("Caption Scoring")
+uploaded_file = st.file_uploader("Choose an image...", type=["png", "jpg"])
+text_input = st.text_input("Type a caption")
+if uploaded_file is not None and text_input:
+    local_image_path = None
+    try:
+        local_image_path = save_file_to_disk(uploaded_file)
+        score = run_inference(local_image_path, text_input, model, tokenizer).tolist()
+        st.image(
+            uploaded_file,
+            caption=text_input,
+            width=None,
+            use_column_width=None,
+            clamp=False,
+            channels="RGB",
+            output_format="auto",
+        )
+        st.write(f"## Score: {score:.2f}")
+    finally:
+        if local_image_path:
+            os.remove(local_image_path)

requirements.txt

@@ -0,0 +1,8 @@
+flax==0.3.4
+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