From 0bdeabbbe4e7f3b1f6b1a5aaf0e6e6eb87d54e12 Mon Sep 17 00:00:00 2001
From: jmoller <jmoller@ginkgobioworks.com>
Date: Thu, 12 Dec 2024 15:35:24 +0000
Subject: [PATCH] Adding antibody support to API

---
 ginkgo_ai_client/queries.py |  1 +
 test/test_models.py         | 14 ++++++++++----
 2 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/ginkgo_ai_client/queries.py b/ginkgo_ai_client/queries.py
index 9fda36e..be4d5cc 100644
--- 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
index 44239f7..e496f49 100644
--- 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,