Update README.md

README.md

@@ -3083,7 +3083,7 @@ widget:
 
 # Model description
 
-This is a LogisticRegressionCV model trained on averages of patch embeddings from the Imagenette dataset..
+This is a LogisticRegressionCV model trained on averages of patch embeddings from the Imagenette dataset. This forms the GAM of an [Emb-GAM](https://arxiv.org/abs/2209.11799) extended to images. Patch embeddings are meant to be extracted with the [`facebook/dino-vitb16` DINO checkpoint](https://huggingface.co/facebook/dino-vitb16).
 
 ## Intended uses & limitations
 
@@ -3145,9 +3145,40 @@ Use the code below to get started with the model.
 <summary> Click to expand </summary>
 
 ```python
-import pickle 
-with open('model.pkl', 'rb') as file: 
-    clf = pickle.load(file)
+from PIL import Image
+from skops import hub_utils
+import torch
+from transformers import ViTFeatureExtractor, ViTModel
+import pickle
+import os
+
+# load DINO
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+feature_extractor = ViTFeatureExtractor.from_pretrained('facebook/dino-vitb16')
+model = ViTModel.from_pretrained('facebook/dino-vitb16').eval().to(device)
+
+# load logistic regression
+os.mkdir('logistic regression')
+hub_utils.download(repo_id='Ramos-Ramos/emb-gam-dino', dst='emb-gam-dino')
+
+with open('emb-gam-dino/model.pkl', 'rb') as file: 
+  logistic_regression = pickle.load(file)
+    
+# load image
+img = Image.open('examples/english_springer.png')
+
+# preprocess image
+inputs = {k: v.to(device) for k, v in feature_extractor(img, return_tensors='pt').items()}
+
+# extract patch embeddings
+with torch.no_grad():
+  patch_embeddings = model(**inputs).last_hidden_state[0, 1:].cpu()
+
+# classify
+pred = logistic_regression.predict(patch_embeddings.mean(dim=0).view(1, -1))
+
+# get patch contributions
+patch_contributions = logistic_regression.coef_ @ patch_embeddings.T.numpy()
 ```
 
 </details>
@@ -3168,11 +3199,16 @@ You can contact the model card authors through following channels:
 
 # Citation
 
-Below you can find information related to citation.
+Below you can find information related to citation. Note that this is **not our own paper**.
 
 **BibTeX:**
 ```
-[More Information Needed]
+@article{singh2022emb,
+  title={Emb-GAM: an Interpretable and Efficient Predictor using Pre-trained Language Models},
+  author={Singh, Chandan and Gao, Jianfeng},
+  journal={arXiv preprint arXiv:2209.11799},
+  year={2022}
+}
 ```