|
3 | 3 | import pytest |
4 | 4 | import torch |
5 | 5 |
|
6 | | -from esm.sdk import client, esmc_client # pyright: ignore |
| 6 | +from esm.sdk import client # pyright: ignore |
7 | 7 | from esm.sdk.api import ( # pyright: ignore |
8 | 8 | ESMProtein, |
9 | 9 | ESMProteinTensor, |
@@ -58,19 +58,19 @@ def test_oss_esmc_client(): |
58 | 58 |
|
59 | 59 | sequence = "MALWMRLLPLLALLALAVPDPAAA" |
60 | 60 | model = "esmc-300m-2024-12" |
61 | | - esmc = esmc_client(model=model, url=URL, token=API_TOKEN) |
| 61 | + esmc_client = client(model=model, url=URL, token=API_TOKEN) |
62 | 62 |
|
63 | 63 | protein = ESMProtein(sequence) |
64 | | - encoded_protein = esmc.encode(input=protein) |
| 64 | + encoded_protein = esmc_client.encode(input=protein) |
65 | 65 | assert isinstance(encoded_protein, ESMProteinTensor) |
66 | 66 |
|
67 | | - decoded_protein = esmc.decode(input=encoded_protein) |
| 67 | + decoded_protein = esmc_client.decode(input=encoded_protein) |
68 | 68 | assert isinstance(decoded_protein, ESMProtein) |
69 | 69 |
|
70 | 70 | logits_config = LogitsConfig( |
71 | 71 | sequence=True, return_embeddings=True, return_hidden_states=True |
72 | 72 | ) |
73 | | - result = esmc.logits(input=encoded_protein, config=logits_config) |
| 73 | + result = esmc_client.logits(input=encoded_protein, config=logits_config) |
74 | 74 | assert isinstance(result, LogitsOutput) |
75 | 75 | assert result.logits is not None |
76 | 76 | assert isinstance(result.logits.sequence, torch.Tensor) |
|
0 commit comments