Update README.md

README.md

@@ -28,17 +28,18 @@ import torch
 import torch.nn.functional as F
 from urllib.request import urlopen
 from PIL import Image
-from open_clip import create_model_from_pretrained, get_tokenizer
+from open_clip import create_model_from_pretrained, get_tokenizer # works on open-clip-torch>=2.23.0, timm>=0.9.8
 
-model, preprocess = create_model_from_pretrained('hf-hub:ViT-SO400M-14-SigLIP')
-tokenizer = get_tokenizer('hf-hub:ViT-SO400M-14-SigLIP')
+model, preprocess = create_model_from_pretrained('hf-hub:timm/ViT-SO400M-14-SigLIP')
+tokenizer = get_tokenizer('hf-hub:timm/ViT-SO400M-14-SigLIP')
 
 image = Image.open(urlopen(
     'https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/beignets-task-guide.png'
 ))
 image = preprocess(image).unsqueeze(0)
 
-text = tokenizer(["a diagram", "a dog", "a cat", "a beignet"], context_length=model.context_length)
+labels_list = ["a dog", "a cat", "a donut", "a beignet"]
+text = tokenizer(labels_list, context_length=model.context_length)
 
 with torch.no_grad(), torch.cuda.amp.autocast():
     image_features = model.encode_image(image)
@@ -46,9 +47,10 @@ with torch.no_grad(), torch.cuda.amp.autocast():
     image_features = F.normalize(image_features, dim=-1)
     text_features = F.normalize(text_features, dim=-1)
 
-    text_probs = (100.0 * image_features @ text_features.T).softmax(dim=-1)
+    text_probs = torch.sigmoid(image_features @ text_features.T * model.logit_scale.exp() + model.logit_bias)
 
-print("Label probs:", text_probs)  # prints: [[0., 0., 0., 1.0]]
+zipped_list = list(zip(labels_list, [round(p.item(), 3) for p in text_probs[0]]))
+print("Label probabilities: ", zipped_list)
 ```
 
 ### With `timm` (for image embeddings)