Model list improvement

app.py

@@ -4,6 +4,7 @@ import image2text
 import intro
 import text2image
 import text2patch
+from config import MODEL_LIST
 
 PAGES = {
     "Introduction": intro,
@@ -13,7 +14,7 @@ PAGES = {
 }
 
 st.sidebar.title("Navigation")
-model = st.sidebar.selectbox("Choose a model", ["koclip-base", "koclip-large"])
+model = st.sidebar.selectbox("Choose a model", MODEL_LIST)
 page = st.sidebar.selectbox("Navigate to...", list(PAGES.keys()))
 
 PAGES[page].app(model)

config.py

@@ -0,0 +1 @@
+MODEL_LIST = ["koclip-base", "koclip-large"]

executables/embed_captions.py

@@ -1,15 +1,18 @@
 import argparse
 import csv
+import numpy as np
+import os
 
+from config import MODEL_LIST
 from utils import load_model
 
 
 def main(args):
     caption_txt_path = args.text_path
     f = open(caption_txt_path)
-    captions = [sent.strip() for sent in f.readlines()
+    captions = [sent.strip() for sent in f.readlines()]
 
-    for model_name in ["koclip-base", "koclip-large"]:
+    for model_name in MODEL_LIST:
         model, processor = load_model(f"koclip/{model_name}")
         captions_processed = [processor(sent,images=None,return_tensors='jax') for sent in captions]
         vec = [np.asarray(model.get_text_features(**c)) for c in captions_processed]

executables/embed_images.py

@@ -7,6 +7,7 @@ from jax import jit
 from PIL import Image
 from tqdm import tqdm
 
+from config import MODEL_LIST
 from utils import load_model
 
 
@@ -15,7 +16,7 @@ def main(args):
     files = list(os.listdir(root))
     for f in files:
         assert f[-4:] == ".jpg"
-    for model_name in ["koclip-base", "koclip-large"]:
+    for model_name in MODEL_LIST:
         model, processor = load_model(f"koclip/{model_name}")
         with tqdm(total=len(files)) as pbar:
             for counter in range(0, len(files), args.batch_size):

utils.py

@@ -3,6 +3,7 @@ import numpy as np
 import streamlit as st
 from transformers import AutoTokenizer, CLIPProcessor, ViTFeatureExtractor
 
+from config import MODEL_LIST
 from koclip import FlaxHybridCLIP
 
 
@@ -25,7 +26,7 @@ def load_index(img_file):
 
 @st.cache(allow_output_mutation=True)
 def load_model(model_name="koclip/koclip-base"):
-    assert model_name in {"koclip/koclip-base", "koclip/koclip-large"}
+    assert model_name in {f"koclip/{model}" for model in MODEL_LIST}
     model = FlaxHybridCLIP.from_pretrained(model_name)
     processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
     processor.tokenizer = AutoTokenizer.from_pretrained("klue/roberta-large")