Update README.md

README.md

@@ -40,10 +40,9 @@ import torch
 # Import the tokenizer and the model
 tokenizer = AutoTokenizer.from_pretrained("InstaDeepAI/nucleotide-transformer-500m-1000g")
 model = AutoModelForMaskedLM.from_pretrained("InstaDeepAI/nucleotide-transformer-500m-1000g")
-
 # Create a dummy dna sequence and tokenize it
-sequences = ['ATTCTG' * 9]
-tokens_ids = tokenizer.batch_encode_plus(sequences, return_tensors="pt")["input_ids"]
+sequences = ["ATTCCGATTCCGATTCCG", "ATTTCTCTCTCTCTCTGAGATCGATCGATCGAT"]
+tokens_ids = tokenizer.batch_encode_plus(sequences, return_tensors="pt", padding="max_length", max_length = max_length)["input_ids"]
 
 # Compute the embeddings
 attention_mask = tokens_ids != tokenizer.pad_token_id
@@ -59,8 +58,11 @@ embeddings = torch_outs['hidden_states'][-1].detach().numpy()
 print(f"Embeddings shape: {embeddings.shape}")
 print(f"Embeddings per token: {embeddings}")
 
+# Add embed dimension axis
+attention_mask = torch.unsqueeze(attention_mask, dim=-1)
+
 # Compute mean embeddings per sequence
-mean_sequence_embeddings = torch.sum(attention_mask.unsqueeze(-1)*embeddings, axis=-2)/torch.sum(attention_mask, axis=-1)
+mean_sequence_embeddings = torch.sum(attention_mask*embeddings, axis=-2)/torch.sum(attention_mask, axis=1)
 print(f"Mean sequence embeddings: {mean_sequence_embeddings}")
 ```