Update README.md

README.md

@@ -4,13 +4,16 @@ library_name: transformers
 datasets:
 - CCDS
 - Ensembl
-pipeline_tag: fill-mask
+pipeline_tag: feature-extraction
 tags:
 - protein language model
 - biology
 widget:
-- text: ( Z [MASK] V L P Y G D E K L S P Y G D G G D V G Q I F s C B L Q D T N N F F G A g Q N K % O P K L G Q I G % S K % u u i e d d R i d D V L k n ( T D K @ p p ^ v
-  example_title: Fill codon mask (E)
+- text: >-
+    ( Z E V L P Y G D E K L S P Y G D G G D V G Q I F s C B L Q D T N N F F G A
+    g Q N K % O P K L G Q I G % S K % u u i e d d R i d D V L k n ( T D K @ p p
+    ^ v
+  example_title: Feature extraction
 ---
 
 # cdsBERT
@@ -29,7 +32,7 @@ import torch
 import torch.nn.functional as F
 from transformers import BertForMaskedLM, BertTokenizer
 
-model = BertForMaskedLM.from_pretrained('lhallee/cdsBERT') # load model
+model = BertModel.from_pretrained('lhallee/cdsBERT') # load model
 tokenizer = BertTokenizer.from_pretrained('lhallee/cdsBERT') # load tokenizer
 device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') # gather device
 model.to(device) # move to device
@@ -59,5 +62,4 @@ The [Gleghorn lab](https://www.gleghornlab.com/) is an interdisciplinary researc
 	doi = {10.1101/2023.09.15.558027},
 	publisher = {Cold Spring Harbor Laboratory},
 	journal = {bioRxiv}
-}
-
+}