Add files

.gitattributes

@@ -32,3 +32,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+cc3m_embeddings.pkl filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,7 +1,7 @@
 ---
 title: FROMAGe
 emoji: 🧀
-sdk: docker
+sdk: gradio
 colorFrom: blue
 colorTo: red
 pinned: true

app.py

@@ -13,10 +13,9 @@ import tempfile
 class FromageChatBot:
     def __init__(self):
         # Download model from HF Hub.
-        huggingface_hub.hf_hub_download(repo_id='jykoh/fromage', filename='pretrained_ckpt.pth.tar')
-        huggingface_hub.hf_hub_download(repo_id='jykoh/fromage', filename='model_args.json')
-        huggingface_hub.hf_hub_download(repo_id='jykoh/fromage', filename='cc3m_embeddings.pkl')
-        self.model = models.load_fromage('./')
+        ckpt_path = huggingface_hub.hf_hub_download(repo_id='jykoh/fromage', filename='pretrained_ckpt.pth.tar')
+        args_path = huggingface_hub.hf_hub_download(repo_id='jykoh/fromage', filename='model_args.json')
+        self.model = models.load_fromage('./', args_path, ckpt_path)
         self.chat_history = ''
         self.input_image = None
 

cc3m_embeddings.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a20fa8168bd72e848ff088820b767383dded455a57ac5dd2d97d43e600402195
+size 2979901225

fromage/models.py

@@ -594,7 +594,7 @@ class Fromage(nn.Module):
     return return_outputs
 
 
-def load_fromage(model_dir: str) -> Fromage:
+def load_fromage(embeddings_dir: str, args_path: str, ckpt_path: str) -> Fromage:
   model_args_path = os.path.join(model_dir, 'model_args.json')
   model_ckpt_path = os.path.join(model_dir, 'pretrained_ckpt.pth.tar')
   embs_paths = [s for s in glob.glob(os.path.join(model_dir, 'cc3m_embeddings*.pkl'))]