Fixed bug

wfondrie · May 10, 2024 · 931c37c · 931c37c
1 parent 7d022f3
commit 931c37c
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 1 deletion.
diff --git a/depthcharge/tokenizers/peptides.py b/depthcharge/tokenizers/peptides.py
@@ -194,7 +194,7 @@ def detokenize(
             tokens=tokens,
             join=join,
             trim_start_token=trim_start_token,
-            trim_stop_token=trim_start_token,
+            trim_stop_token=trim_stop_token,
         )
 
         if self.reverse:

diff --git a/tests/unit_tests/test_tokenizers/test_peptides.py b/tests/unit_tests/test_tokenizers/test_peptides.py
@@ -115,3 +115,25 @@ def test_almost_compliant_proform():
     """Test initializing with a peptide without an expicit mass sign."""
     tokenizer = PeptideTokenizer.from_proforma("[10]-EDITHR")
     assert "[+10.000000]-" in tokenizer.residues
+
+
+@pytest.mark.parametrize(
+    ("start", "stop", "expected"),
+    [
+        (True, True, "ACD"),
+        (True, False, "ACD$E"),
+        (False, True, "?ACD"),
+        (False, False, "?ACD$E"),
+    ],
+)
+def test_trim(start, stop, expected):
+    """Test that the start and stop tokens can be trimmed."""
+    tokenizer = PeptideTokenizer(start_token="?")
+    tokens = torch.tensor([[2, 3, 4, 5, 1, 6]])
+    out = tokenizer.detokenize(
+        tokens,
+        trim_start_token=start,
+        trim_stop_token=stop,
+    )
+
+    assert out[0] == expected