Refactor test_on_image.py

test_on_image.py

@@ -8,10 +8,8 @@ from transformers import AutoTokenizer
 from modeling_hybrid_clip import FlaxHybridCLIP
 from run_hybrid_clip import Transform
 
-model = FlaxHybridCLIP.from_pretrained("clip_spanish_1_percent")
-tokenizer = AutoTokenizer.from_pretrained("dccuchile/bert-base-spanish-wwm-cased")
 
-def prepare_image(image_path):
+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)
@@ -19,18 +17,23 @@ def prepare_image(image_path):
     pixel_values = torch.stack([preprocessed_image]).permute(0, 2, 3, 1).numpy()
     return pixel_values
 
-def prepare_text(text):
+def prepare_text(text, tokenizer):
     return tokenizer(text, return_tensors="np")
 
-def run_inference(image_path, text):
-    pixel_values = prepare_image(image_path)
-    input_text = prepare_text(text)
+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"], token_type_ids=input_text["token_type_ids"], train=False, return_dict=True)
     logits = model_output["logits_per_image"]
-    score = jax.nn.sigmoid(logits)
+    score = jax.nn.sigmoid(logits)[0][0]
     return score
 
-image_path = f"/home/{os.environ['USER']}/data/wit_scale_converted/Self_Portrait_by_David_Allan.jpg"
-text = "Patio interior de un edificio"
 
-print(run_inference(image_path, text))
+if __name__ == "__main__":
+    model = FlaxHybridCLIP.from_pretrained("clip_spanish_141230_samples")
+    tokenizer = AutoTokenizer.from_pretrained("dccuchile/bert-base-spanish-wwm-cased")
+
+    image_path = f"/home/{os.environ['USER']}/data/wit_scale_converted/Santuar.jpg"
+    text = "Fachada del Santuario"
+
+    print(run_inference(image_path, text, model, tokenizer))