Update README.md

README.md

@@ -11,7 +11,41 @@ tags:
 # Note:
 This model is copied version of mRNABERT which removes the FlashAttention integration with Trition. This allows the model to be installed off HuggingFace without having to uninstall Triton. Running the below example code yields identical output compared to the original verison.
 ```
-# TODO
+import torch
+from transformers import AutoTokenizer, AutoModel
+from transformers.models.bert.configuration_bert import BertConfig
+
+config = BertConfig.from_pretrained("Taykhoom/mRNABERT-no-flashattention")
+tokenizer = AutoTokenizer.from_pretrained("Taykhoom/mRNABERT-no-flashattention")
+model = AutoModel.from_pretrained("Taykhoom/mRNABERT-no-flashattention", trust_remote_code=True, config=config)
+
+seq = ["A T C G G A GGG CCC TTT", 
+       "A T C G", 
+       "TTT CCC GAC ATG"]  #Separate the sequences with spaces.
+
+encoding = tokenizer.batch_encode_plus(seq, add_special_tokens=True, padding='longest', return_tensors="pt")
+
+input_ids = encoding['input_ids']
+attention_mask = encoding['attention_mask'] 
+
+output = model(input_ids=input_ids, attention_mask=attention_mask)
+last_hidden_state = output[0]
+
+attention_mask = attention_mask.unsqueeze(-1).expand_as(last_hidden_state)  # Shape : [batch_size, seq_length, hidden_size]
+
+# Sum embeddings along the batch dimension
+sum_embeddings = torch.sum(last_hidden_state * attention_mask, dim=1)  
+
+# Also sum the masks along the batch dimension
+sum_masks = attention_mask.sum(1)  
+
+# Compute mean embedding.
+mean_embedding = sum_embeddings / sum_masks  #Shape:[batch_size, hidden_size]
+
+print(torch.mean(mean_embedding, dim=1))
+# Should output: tensor([-0.0209, -0.0156, -0.0201], device='cuda:0', grad_fn=<MeanBackward1>)
+# this is the same as the original version of mRNABERT (checked using original installation instructions)
+
 ```
 
 # Original README: