Return embeddings only when specified

Merck · Jan 17, 2022 · d250781 · d250781
1 parent d8eeb99
commit d250781
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/biophi/humanization/methods/sapiens/roberta.py b/biophi/humanization/methods/sapiens/roberta.py
@@ -73,7 +73,10 @@ def predict_proba(self, seq, remove_special=True, return_all_hiddens=False):
  pred = pd.DataFrame(pred.numpy(), columns=self.interface.task.target_dictionary.symbols)
  if remove_special:
  pred.drop(['<s>', '<pad>', '</s>', '<unk>', '<mask>'], axis=1, inplace=True)
- return pred, extra
+ if return_all_hiddens:
+ return pred, extra
+ else:
+ return pred
 
  def _is_adding_bos(self):
  if isinstance(self.interface.task, SentencePredictionTask):