added sub_prompting

app.py

@@ -5,9 +5,9 @@ import time
 import json
 sys.path.append("./virtex/")
 
-def gen_show_caption():
+def gen_show_caption(sub_prompt=None):
     with st.spinner("Generating Caption"):
-        subreddit, caption = virtexModel.predict(image_dict)
+        subreddit, caption = virtexModel.predict(image_dict, sub_prompt=sub_prompt)
         st.header("Predicted Caption:\n\n")
         st.subheader(f"r/{subreddit}:\t{caption}\n")
     
@@ -68,7 +68,7 @@ else:
     show = st.image(image)
     show.image(image, "Your Image")
     
-    gen_show_caption()
+    gen_show_caption(sub)
         
     image.close()
 

model.py

@@ -54,7 +54,14 @@ class VirTexModel():
         if sub_prompt is None:
             subreddit_tokens = torch.tensor([self.model.sos_index], device=self.device).long()
         else:
-            subreddit_tokens = torch.tensor([self.tokenizer.token_to_id(sub_prompt)], device=self.device).long()
+            subreddit_tokens = " ".join(ws.segment(ws.clean(sub_prompt)))
+            subreddit_tokens = (
+                [self.model.sos_index] + 
+                self.tokenizer.encode(subreddit_tokens) +
+                [tokenizer.token_to_id("[SEP]")]
+                               )
+            
+            
         predictions: List[Dict[str, Any]] = []
         
         is_valid_subreddit = False