Update README.md

README.md

@@ -74,17 +74,17 @@ Use the code below to get started with the model.
 from PIL import Image
 from skops import hub_utils
 import torch
-from transformers import ViTFeatureExtractor, ViTModel
+from transformers import AutoFeatureExtractor, AutoModel
 import pickle
 import os
 
-# load DINO
+# load embedding model
 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)
+feature_extractor = AutoFeatureExtractor.from_pretrained('facebook/dino-vitb16')
+model = AutoModel.from_pretrained('facebook/dino-vitb16').eval().to(device)
 
 # load logistic regression
-os.mkdir('logistic regression')
+os.mkdir('emb-gam-dino')
 hub_utils.download(repo_id='Ramos-Ramos/emb-gam-dino', dst='emb-gam-dino')
 
 with open('emb-gam-dino/model.pkl', 'rb') as file: 
@@ -101,7 +101,7 @@ 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))
+pred = logistic_regression.predict(patch_embeddings.sum(dim=0, keepdim=True))
 
 # get patch contributions
 patch_contributions = logistic_regression.coef_ @ patch_embeddings.T.numpy()