User BERTIN model

app.py

@@ -1,6 +1,7 @@
 import os
 import sys
 
+import jax
 import streamlit as st
 import transformers
 from huggingface_hub import snapshot_download
@@ -10,7 +11,7 @@ LOCAL_PATH = snapshot_download("flax-community/clip-spanish")
 sys.path.append(LOCAL_PATH)
 
 from modeling_hybrid_clip import FlaxHybridCLIP
-from test_on_image import run_inference
+from test_on_image import prepare_image, prepare_text
 
 
 def save_file_to_disk(uplaoded_file):
@@ -22,16 +23,30 @@ def save_file_to_disk(uplaoded_file):
 
 @st.cache(
     hash_funcs={
-        transformers.models.bert.tokenization_bert_fast.BertTokenizerFast: id,
+        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("dccuchile/bert-base-spanish-wwm-cased")
+    tokenizer = AutoTokenizer.from_pretrained("bertin-project/bertin-roberta-base-spanish")
     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()