Update README.md

README.md

@@ -18,13 +18,9 @@ logger = logging.getLogger(__name__)
 
 
 def tokenize_protein(example, protein_tokenizer=None, padding=None):
-    # check https://github.com/huggingface/transformers/blob/41aef33758ae166291d72bc381477f2db84159cf/src/transformers/models/esm/tokenization_esm.py#L100
     protein_seqs = example["prot_seq"]
     
-    protein_inputs = protein_tokenizer(
-                        protein_seqs, padding=padding, 
-                        add_special_tokens=True, # default is True, no need to add cls and eos manually
-                    ) # results in <cls> + seq + <eos> (no <sep> for ESM)
+    protein_inputs = protein_tokenizer(protein_seqs, padding=padding, add_special_tokens=True)
     example["protein_input_ids"] = protein_inputs.input_ids
     example["protein_attention_mask"] = protein_inputs.attention_mask
     
@@ -42,18 +38,7 @@ def label_embedding(labels, text_tokenizer, text_model, device):
             label_input_ids = torch.tensor(label_input_ids, dtype=torch.long, device=device).unsqueeze(0)
             attention_mask = label_input_ids != text_tokenizer.pad_token_id
 
-            text_outputs = text_model(
-                label_input_ids,
-                attention_mask=attention_mask,
-                position_ids=None,
-                head_mask=None,
-                inputs_embeds=None,
-                encoder_hidden_states=None,
-                encoder_attention_mask=None,
-                output_attentions=None,
-                output_hidden_states=None,
-                return_dict=None,
-            )
+            text_outputs = text_model(label_input_ids, attention_mask=attention_mask)
 
             label_feature.append(text_outputs["text_feature"])
     label_feature = torch.cat(label_feature, dim=0)
@@ -75,11 +60,7 @@ def zero_shot_eval(logger, device,
             protein_input_ids = torch.tensor(data["protein_input_ids"], dtype=torch.long, device=device).unsqueeze(0)
             attention_mask = torch.tensor(data["protein_attention_mask"], dtype=torch.long, device=device).unsqueeze(0)
             
-            protein_outputs = protein_model(
-                protein_input_ids,
-                attention_mask=attention_mask,
-                position_ids=None, # it's ok to set `position_ids`` as None: https://github.com/huggingface/transformers/blob/41aef33758ae166291d72bc381477f2db84159cf/src/transformers/models/esm/modeling_esm.py#L195
-            )
+            protein_outputs = protein_model(protein_input_ids, attention_mask=attention_mask)
 
             protein_feature = protein_outputs["protein_feature"]
             protein_feature = protein_feature / protein_feature.norm(dim=-1, keepdim=True)