add tokenizers

app.py

@@ -6,19 +6,17 @@ import io
 import time
 
 @st.cache(show_spinner=True,allow_output_mutation=True)
-def load_model(model_name):
-    if model_name.startswith('bert'):
-        from transformers import BertTokenizer
-        tokenizer = BertTokenizer.from_pretrained(model_name)
-    elif model_name.startswith('gpt2'):
-        from transformers import GPT2Tokenizer
-        tokenizer = GPT2Tokenizer.from_pretrained(model_name)
-    elif model_name.startswith('roberta'):
-        from transformers import RobertaTokenizer
-        tokenizer = RobertaTokenizer.from_pretrained(model_name)
-    elif model_name.startswith('albert'):
-        from transformers import AlbertTokenizer
-        tokenizer = AlbertTokenizer.from_pretrained(model_name)
+def load_model(tokenizer_name):
+    from transformers import AutoTokenizer
+    model_name_dict = {
+        "BERT":"bert-base-uncased",
+        "RoBERTa":"roberta-base",
+        "ALBERT":"albert-v2-base",
+        "GPT2":"gpt2",
+        "Llama":"meta-lama/Llama-2-7b-chat-hf",
+        "Gemma":"google/gemma-7b",
+        }
+    tokenizer = AutoTokenizer.from_pretrained(model_name_dict[tokenizer_name])
     return tokenizer
 
 def generate_markdown(text,color='black',font='Arial',size=20):
@@ -26,10 +24,11 @@ def generate_markdown(text,color='black',font='Arial',size=20):
 
 def TokenizeText(sentence,tokenizer_name):
     if len(sentence)>0:
-        if tokenizer_name.startswith('gpt2'):
-            input_sent = tokenizer(sentence)['input_ids']
-        else:
-            input_sent = tokenizer(sentence)['input_ids'][1:-1]
+        #if tokenizer_name.startswith('gpt2'):
+        #    input_sent = tokenizer(sentence)['input_ids']
+        #else:
+        #    input_sent = tokenizer(sentence)['input_ids'][1:-1]
+        input_sent = tokenizer(sentence)['input_ids']
         encoded_sent = [str(token) for token in input_sent]
         decoded_sent = [tokenizer.decode([token]) for token in input_sent]
         num_tokens = len(decoded_sent)
@@ -100,10 +99,8 @@ if __name__=='__main__':
     # Select and load the tokenizer
     st.sidebar.write('1. Choose the tokenizer from below')
     tokenizer_name = st.sidebar.selectbox('',
-                                            ('bert-base-uncased','bert-large-cased',
-                                            'gpt2','gpt2-large',
-                                            'roberta-base','roberta-large',
-                                            'albert-base-v2','albert-xxlarge-v2'),index=7)
+                                            ("BERT","RoBERTa","ALBERT",
+                                             "GPT2","Llama","Gemma"))
     tokenizer = load_model(tokenizer_name)
 
     st.sidebar.write('2. Optional settings')
@@ -135,10 +132,11 @@ if __name__=='__main__':
 
     else:
         if detokenize:
-            if tokenizer_name.startswith('gpt2'):
-                default_tokens = tokenizer('Tokenizers decompose bigger words into smaller tokens')['input_ids']
-            else:
-                default_tokens = tokenizer('Tokenizers decompose bigger words into smaller tokens')['input_ids'][1:-1]
+            #if tokenizer_name.startswith('gpt2'):
+            #    default_tokens = tokenizer('Tokenizers decompose bigger words into smaller tokens')['input_ids']
+            #else:
+            #    default_tokens = tokenizer('Tokenizers decompose bigger words into smaller tokens')['input_ids'][1:-1]
+            default_tokens = tokenizer('Tokenizers decompose bigger words into smaller tokens')['input_ids']
             sentence = st.text_input(f'Tokenized IDs',value=' '.join([str(token) for token in default_tokens]))
             num_tokens = DeTokenizeText(sentence)
         else: