()
| 14 | os.environ["CUDA_VISIBLE_DEVICES"] = "" |
| 15 | |
| 16 | def test_inference(): |
| 17 | model, _, preprocess = create_model_and_transforms('ViT-B-32-quickgelu', pretrained='laion400m_e32') |
| 18 | |
| 19 | current_dir = os.path.dirname(os.path.realpath(__file__)) |
| 20 | |
| 21 | image = preprocess(Image.open(current_dir + "/../docs/CLIP.png")).unsqueeze(0) |
| 22 | text = tokenizer.tokenize(["a diagram", "a dog", "a cat"]) |
| 23 | |
| 24 | with torch.no_grad(): |
| 25 | image_features = model.encode_image(image) |
| 26 | text_features = model.encode_text(text) |
| 27 | |
| 28 | text_probs = (100.0 * image_features @ text_features.T).softmax(dim=-1) |
| 29 | |
| 30 | assert text_probs.cpu().numpy()[0].tolist() == [1.0, 0.0, 0.0] |
| 31 | |
| 32 | |
| 33 | def test_metaclip2_inference(): |
nothing calls this directly
no test coverage detected