Adding antibody support to API

ginkgobioworks · Dec 12, 2024 · 0bdeabb · 0bdeabb
1 parent e0a96e8
commit 0bdeabb
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 4 deletions.
diff --git a/ginkgo_ai_client/queries.py b/ginkgo_ai_client/queries.py
@@ -48,6 +48,7 @@ def write_to_jsonl(self, path: str):
     "esm2-3B": "protein",
     "ginkgo-maskedlm-3utr-v1": "dna",
     "lcdna": "nucleotide",
+    "abdiffusion": "protein",
 }
 
 _maskedlm_models_properties_str = "\n".join(

diff --git a/test/test_models.py b/test/test_models.py
@@ -63,12 +63,18 @@ def test_promoter_activity():
     assert "heart" in response.activity_by_tissue
     assert "liver" in response.activity_by_tissue
 
-
-def test_diffusion_masked_inference():
+@pytest.mark.parametrize(
+    "model, sequence, expected_sequence",
+    [
+        ("lcdna", "ATRGAyAtg<mask>TAC<mask>"),
+        ("abdiffusion", "MCL<mask>YAFVATDA<mask>DDT"),
+    ],
+)
+def test_diffusion_masked_inference(model, sequence):
     client = GinkgoAIClient()
     query = DiffusionMaskedQuery(
-        sequence="ATRGAyAtg<mask>TAC<mask>",  #upper and lower cases
-        model="lcdna",
+        sequence=sequence,  #upper and lower cases
+        model=model,
         temperature=0.5,
         decoding_order_strategy="entropy",
         unmaskings_per_step=2,