From 1cf72aa1b6c95d94ef0876c69547430f51d969ab Mon Sep 17 00:00:00 2001 From: Adibvafa Date: Fri, 20 Sep 2024 16:28:59 -0400 Subject: [PATCH] Add support for non-deterministic test with temperature. --- tests/test_CodonPrediction.py | 23 +++++++++++++---------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/tests/test_CodonPrediction.py b/tests/test_CodonPrediction.py index 3617c2b..28cb49e 100644 --- a/tests/test_CodonPrediction.py +++ b/tests/test_CodonPrediction.py @@ -44,7 +44,8 @@ def test_predict_dna_sequence_valid_input(self): def test_predict_dna_sequence_non_deterministic(self): protein_sequence = "MFWY" organism = "Escherichia coli general" - num_iterations = 64 + num_iterations = 100 + temperatures = [0.2, 0.5, 0.8] possible_outputs = set() possible_encodings_wo_stop = { "ATGTTTTGGTAT", @@ -54,15 +55,17 @@ def test_predict_dna_sequence_non_deterministic(self): } for _ in range(num_iterations): - result = predict_dna_sequence( - protein=protein_sequence, - organism=organism, - device=self.device, - tokenizer=self.tokenizer, - model=self.model, - deterministic=False, - ) - possible_outputs.add(result.predicted_dna[:-3]) # Remove stop codon + for temperature in temperatures: + result = predict_dna_sequence( + protein=protein_sequence, + organism=organism, + device=self.device, + tokenizer=self.tokenizer, + model=self.model, + deterministic=False, + temperature=temperature, + ) + possible_outputs.add(result.predicted_dna[:-3]) # Remove stop codon self.assertEqual(possible_outputs, possible_encodings_wo_stop)