Create synonymous_logit_processor.py

synonymous_logit_processor.py

@@ -0,0 +1,42 @@
+class SynonymMaskingLogitsProcessor(LogitsProcessor):
+    def __init__(self, current_aa, tokenizer, aa_to_codon):
+        self.current_aa = current_aa
+        self.tokenizer = tokenizer
+        self.aa_to_codon = aa_to_codon
+
+    def __call__(self, input_ids: torch.Tensor, scores: torch.Tensor) -> torch.Tensor:
+        synonymous_codons = self.aa_to_codon.get(self.current_aa, [])
+        synonym_token_ids = self.tokenizer.convert_tokens_to_ids(synonymous_codons)
+        mask = torch.ones_like(scores) * -float('inf')
+        mask[:, synonym_token_ids] = 0
+        return scores + mask
+
+def generate_candidate_codons_with_generate(initial_codons, temperature=1.0, top_k=None, top_p=None):
+    optimized_codons = []
+    current_sequence_tokens = [tokenizer.bos_token_id]
+
+    for codon in initial_codons:
+        aa = str(Seq(codon).translate())
+        logits_processor = [SynonymMaskingLogitsProcessor(aa, tokenizer, aa_to_codon_human)]
+
+        input_ids = torch.tensor([current_sequence_tokens])#.to(device)
+
+        output = model.generate(
+            input_ids,
+            max_length=len(current_sequence_tokens) + 1,
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p,
+            num_return_sequences=1,
+            pad_token_id=tokenizer.eos_token_id,
+            logits_processor=logits_processor,
+            do_sample=True # Ensure sampling is used for temperature, top_k, top_p
+        )
+
+        next_token_id = output[0][-1].item()
+        predicted_codon = tokenizer.decode([next_token_id])
+
+        optimized_codons.append(predicted_codon.upper())
+        current_sequence_tokens.append(next_token_id)
+
+    return optimized_codons