nucleus size updates in advanced"

app.py

@@ -6,6 +6,11 @@ import json
 sys.path.append("./virtex/")
 from model import *
 
+# # TODO:
+# - Reformat the model introduction
+# - Center the images using the 3 column method
+# - Make the iterative text generation
+
 def gen_show_caption(sub_prompt=None, cap_prompt = ""):
     with st.spinner("Generating Caption"):
         if sub_prompt is None and cap_prompt is not "":
@@ -102,6 +107,7 @@ num_captions=1
 if advanced:
     nuc_size = st.sidebar.slider("Nucelus Size:", min_value=0.0, max_value=1.0, value=0.8, step=0.05)
     num_captions = st.sidebar.select_slider("Number of Captions to Predict", options=[1,2,3,4,5], value=1)
+    virtexModel.model.decoder.nucleus_size = nuc_size
 
 if uploaded_image is None and submitted:
     st.write("Please select a file to upload")
@@ -130,14 +136,4 @@ else:
     show.image(show_image, "Your Image")
 
     for i in range(num_captions):
-        gen_show_caption(sub, imageLoader.text_transform(cap_prompt))
-
-    # from model import *
-    # sample_images = get_samples()
-    # v, il = VirTexModel(), ImageLoader()
-
-    # for s in sample_images:
-    #     subreddit, caption = v.predict(il.load(s))
-    #     print("=====================")
-    #     print(subreddit)
-    #     print(caption)
+        gen_show_caption(sub, imageLoader.text_transform(cap_prompt))

model.py

@@ -61,7 +61,7 @@ class VirTexModel():
         self.device = 'cpu'
         self.tokenizer = TokenizerFactory.from_config(self.config)
         self.model = PretrainingModelFactory.from_config(self.config).to(self.device)
-        CheckpointManager(model=self.model).load("./checkpoint_last5.pth")
+        CheckpointManager(model=self.model).load(MODEL_PATH)
         self.model.eval()
         self.valid_subs = json.load(open(VALID_SUBREDDITS_PATH))