MLM bert added for ZS

.gitignore

@@ -2,5 +2,6 @@ venv/
 #exclude model files as they are large
 sentiment_model_dir/pytorch_model.bin
 zs_model_dir/pytorch_model.bin
+zs_mlm_dir/pytorch_model.bin
 #sent_clf_onnx_dir/
 #zs_onnx_dir/

app.py

@@ -3,6 +3,7 @@ import pandas as pd
 import streamlit as st
 from streamlit_text_rating.st_text_rater import st_text_rater
 from transformers import AutoTokenizer,AutoModelForSequenceClassification
+from transformers import AutoModelForMaskedLM
 import onnxruntime as ort
 import os
 import time
@@ -11,8 +12,15 @@ import plotly.graph_objects as go
 global _plotly_config
 _plotly_config={'displayModeBar': False}
 
-from sentiment_clf_helper import classify_sentiment,create_onnx_model_sentiment,classify_sentiment_onnx
-from zeroshot_clf_helper import zero_shot_classification,create_onnx_model_zs,zero_shot_classification_onnx
+from sentiment_clf_helper import (classify_sentiment,
+                                  create_onnx_model_sentiment,
+                                  classify_sentiment_onnx)
+
+from zeroshot_clf_helper import (zero_shot_classification,
+                                 create_onnx_model_zs_nli,
+                                 create_onnx_model_zs_mlm,
+                                 zero_shot_classification_nli_onnx,
+                                 zero_shot_classification_fillmask_onnx)
 
 import multiprocessing
 total_threads=multiprocessing.cpu_count()#for ort inference
@@ -36,6 +44,10 @@ zs_onnx_mdl_dir=config['ZEROSHOT_CLF']['zs_onnx_mdl_dir']
 zs_onnx_mdl_name=config['ZEROSHOT_CLF']['zs_onnx_mdl_name']
 zs_onnx_quant_mdl_name=config['ZEROSHOT_CLF']['zs_onnx_quant_mdl_name']
 
+zs_mlm_chkpt=config['ZEROSHOT_MLM']['zs_mlm_chkpt']
+zs_mlm_mdl_dir=config['ZEROSHOT_MLM']['zs_mlm_mdl_dir']
+zs_mlm_onnx_mdl_dir=config['ZEROSHOT_MLM']['zs_mlm_onnx_mdl_dir']
+zs_mlm_onnx_mdl_name=config['ZEROSHOT_MLM']['zs_mlm_onnx_mdl_name']
 
 st.set_page_config(  # Alternate names: setup_page, page, layout
     layout="wide",  # Can be "centered" or "wide". In the future also "dashboard", etc.
@@ -43,7 +55,6 @@ st.set_page_config(  # Alternate names: setup_page, page, layout
     page_title='None',  # String or None. Strings get appended with "• Streamlit".
 )
 
-
 padding_top = 0
 st.markdown(f"""
     <style>
@@ -98,16 +109,24 @@ session_options_ort.inter_op_num_threads=1
 # session_options_ort.execution_mode = session_options_ort.ExecutionMode.ORT_SEQUENTIAL
 
 @st.cache(allow_output_mutation=True, suppress_st_warning=True, max_entries=None, ttl=None)
-def create_model_dir(chkpt, model_dir):
+def create_model_dir(chkpt, model_dir,task_type):
     if not os.path.exists(model_dir):
         try:
             os.mkdir(path=model_dir)
         except:
             pass
-        _model = AutoModelForSequenceClassification.from_pretrained(chkpt)
-        _tokenizer = AutoTokenizer.from_pretrained(chkpt)
-        _model.save_pretrained(model_dir)
-        _tokenizer.save_pretrained(model_dir)
+        if task_type=='classification':
+            _model = AutoModelForSequenceClassification.from_pretrained(chkpt)
+            _tokenizer = AutoTokenizer.from_pretrained(chkpt)
+            _model.save_pretrained(model_dir)
+            _tokenizer.save_pretrained(model_dir)
+        elif task_type=='mlm':
+            _model=AutoModelForMaskedLM.from_pretrained(chkpt)
+            _tokenizer=AutoTokenizer.from_pretrained(chkpt)
+            _model.save_pretrained(model_dir)
+            _tokenizer.save_pretrained(model_dir)
+        else:
+            pass
     else:
         pass
 
@@ -125,7 +144,7 @@ with st.sidebar:
 ############### Pre-Download & instantiate objects for sentiment analysis *********************** START **********************
 
 # #create model/token dir for sentiment classification for faster inference
-create_model_dir(chkpt=sent_chkpt, model_dir=sent_mdl_dir)
+create_model_dir(chkpt=sent_chkpt, model_dir=sent_mdl_dir,task_type='classification')
 
 
 @st.cache(allow_output_mutation=True, suppress_st_warning=True, max_entries=None, ttl=None)
@@ -135,7 +154,7 @@ def sentiment_task_selected(task,
                             sent_onnx_mdl_dir=sent_onnx_mdl_dir,
                             sent_onnx_mdl_name=sent_onnx_mdl_name,
                             sent_onnx_quant_mdl_name=sent_onnx_quant_mdl_name):
-    #model & tokenizer initialization for normal sentiment classification
+    ##model & tokenizer initialization for normal sentiment classification
     # model_sentiment=AutoModelForSequenceClassification.from_pretrained(sent_chkpt)
     # tokenizer_sentiment=AutoTokenizer.from_pretrained(sent_chkpt)
     tokenizer_sentiment = AutoTokenizer.from_pretrained(sent_mdl_dir)
@@ -152,18 +171,17 @@ def sentiment_task_selected(task,
 ############## Pre-Download & instantiate objects for sentiment analysis ********************* END **********************************
 
 
-############### Pre-Download & instantiate objects for Zero shot clf *********************** START **********************
+############### Pre-Download & instantiate objects for Zero shot clf NLI *********************** START **********************
 
 # create model/token dir for zeroshot clf -- already created so not required
-create_model_dir(chkpt=zs_chkpt, model_dir=zs_mdl_dir)
+create_model_dir(chkpt=zs_chkpt, model_dir=zs_mdl_dir,task_type='classification')
 
 @st.cache(allow_output_mutation=True, suppress_st_warning=True, max_entries=None, ttl=None)
-def zs_task_selected(task,
-                     zs_chkpt=zs_chkpt ,
-                     zs_mdl_dir=zs_mdl_dir,
-                     zs_onnx_mdl_dir=zs_onnx_mdl_dir,
-                     zs_onnx_mdl_name=zs_onnx_mdl_name,
-                     zs_onnx_quant_mdl_name=zs_onnx_quant_mdl_name):
+def zs_nli_task_selected(task,
+                     zs_chkpt ,
+                     zs_mdl_dir,
+                     zs_onnx_mdl_dir,
+                     zs_onnx_mdl_name):
 
     ##model & tokenizer initialization for normal ZS classification
     # model_zs=AutoModelForSequenceClassification.from_pretrained(zs_chkpt)
@@ -171,16 +189,46 @@ def zs_task_selected(task,
     # tokenizer_zs=AutoTokenizer.from_pretrained(zs_chkpt)
     tokenizer_zs = AutoTokenizer.from_pretrained(zs_mdl_dir)
 
-    # # create onnx model for zeroshot but once created locally comment it out.
-    # create_onnx_model_zs()
+    ## create onnx model for zeroshot but once created locally comment it out.
+    #create_onnx_model_zs_nli()
 
     #create inference session from onnx model
     zs_session = ort.InferenceSession(f"{zs_onnx_mdl_dir}/{zs_onnx_mdl_name}",sess_options=session_options_ort)
-    # zs_session_quant = ort.InferenceSession(f"{zs_onnx_mdl_dir}/{zs_onnx_quant_mdl_name}")
 
     return tokenizer_zs,zs_session
 
-############## Pre-Download & instantiate objects for Zero shot analysis ********************* END **********************************
+############## Pre-Download & instantiate objects for Zero shot NLI analysis ********************* END **********************************
+
+
+############### Pre-Download & instantiate objects for Zero shot clf NLI *********************** START **********************
+## create model/token dir for zeroshot clf -- already created so not required
+# create_model_dir(chkpt=zs_mlm_chkpt, model_dir=zs_mlm_mdl_dir, task_type='mlm')
+
+@st.cache(allow_output_mutation=True, suppress_st_warning=True, max_entries=None, ttl=None)
+def zs_mlm_task_selected(task,
+                         zs_mlm_chkpt=zs_mlm_chkpt,
+                         zs_mlm_mdl_dir=zs_mlm_mdl_dir,
+                         zs_mlm_onnx_mdl_dir=zs_mlm_onnx_mdl_dir,
+                         zs_mlm_onnx_mdl_name=zs_mlm_onnx_mdl_name):
+    ##model & tokenizer initialization for normal ZS classification
+    model_zs_mlm=AutoModelForMaskedLM.from_pretrained(zs_mlm_mdl_dir)
+    ##we just need tokenizer for inference and not model since onnx model is already saved
+    # tokenizer_zs_mlm=AutoTokenizer.from_pretrained(zs_mlm_chkpt)
+    tokenizer_zs_mlm = AutoTokenizer.from_pretrained(zs_mlm_mdl_dir)
+
+    # create onnx model for zeroshot but once created locally comment it out.
+    create_onnx_model_zs_mlm(_model=model_zs_mlm,
+                             _tokenizer=tokenizer_zs_mlm,
+                             zs_mlm_onnx_mdl_dir=zs_mlm_onnx_mdl_dir)
+
+    # create inference session from onnx model
+    zs_session_mlm = ort.InferenceSession(f"{zs_mlm_onnx_mdl_dir}/{zs_mlm_onnx_mdl_name}", sess_options=session_options_ort)
+
+    return tokenizer_zs_mlm, zs_session_mlm
+
+
+############## Pre-Download & instantiate objects for Zero shot MLM analysis ********************* END **********************************
+
 
 if select_task=='README':
     st.header("NLP Summary")
@@ -196,7 +244,7 @@ if select_task == 'Detect Sentiment':
     t2 = time.time()
     st.write(f"Total time to load Model is {(t2-t1)*1000:.1f} ms")
 
-    st.header("You are now performing Sentiment Analysis")
+    st.subheader("You are now performing Sentiment Analysis")
     input_texts = st.text_input(label="Input texts separated by comma")
     c1,c2,_,_=st.columns(4)
 
@@ -223,35 +271,73 @@ if select_task == 'Detect Sentiment':
 
 if select_task=='Zero Shot Classification':
     t1=time.time()
-    tokenizer_zs,session_zs = zs_task_selected(task=select_task)
-    # tokenizer_zs= AutoTokenizer.from_pretrained(zs_mdl_dir)
-    # session_zs = ort.InferenceSession(f"{zs_onnx_mdl_dir}/{zs_onnx_mdl_name}")
+    tokenizer_zs,session_zs = zs_nli_task_selected(task=select_task ,
+                                                   zs_chkpt=zs_chkpt,
+                                                   zs_mdl_dir=zs_mdl_dir,
+                                                   zs_onnx_mdl_dir=zs_onnx_mdl_dir,
+                                                   zs_onnx_mdl_name=zs_onnx_mdl_name)
     t2 = time.time()
-    st.write(f"Total time to load Model is {(t2-t1)*1000:.1f} ms")
+    st.write(f"Total time to load NLI Model is {(t2-t1)*1000:.1f} ms")
 
-    st.header("You are now performing Zero Shot Classification")
+    t1=time.time()
+    tokenizer_zs_mlm,session_zs_mlm = zs_mlm_task_selected(task=select_task,
+                                                           zs_mlm_chkpt=zs_mlm_chkpt,
+                                                           zs_mlm_mdl_dir=zs_mlm_mdl_dir,
+                                                           zs_mlm_onnx_mdl_dir=zs_mlm_onnx_mdl_dir,
+                                                           zs_mlm_onnx_mdl_name=zs_mlm_onnx_mdl_name)
+    t2 = time.time()
+    st.write(f"Total time to load MLM Model is {(t2-t1)*1000:.1f} ms")
+
+    st.subheader("Zero Shot Classification using NLI")
     input_texts = st.text_input(label="Input text to classify into topics")
     input_lables = st.text_input(label="Enter labels separated by commas")
+    input_hypothesis = st.text_input(label="Enter your hypothesis",value="This is an example of")
 
-    c1,_,_,_=st.columns(4)
+    c1,c2,_,=st.columns(3)
 
     with c1:
-        response1=st.button("Compute (ONNX runtime)")
+        response1=st.button("Compute using NLI approach (ONNX runtime)")
+
+    with c2:
+        response2=st.button("Compute using Fill-Mask approach(ONNX runtime)")
 
     if response1:
         start = time.time()
-        df_output = zero_shot_classification_onnx(premise=input_texts, labels=input_lables, _session=session_zs,
-                                               _tokenizer=tokenizer_zs)
+        df_output = zero_shot_classification_nli_onnx(premise=input_texts,
+                                                      labels=input_lables,
+                                                      hypothesis=input_hypothesis,
+                                                      _session=session_zs,
+                                                     _tokenizer=tokenizer_zs,
+                                                      )
         end = time.time()
-        st.write("")
         st.write(f"Time taken for computation {(end-start)*1000:.1f} ms")
         fig = px.bar(x='Probability',
                      y='labels',
                      text='Probability',
                      data_frame=df_output,
-                     title='Zero Shot Normalized Probabilities')
+                     title='Zero Shot NLI Normalized Probabilities')
+
+        st.plotly_chart(fig, config=_plotly_config)
+
+    elif response2:
+        start=time.time()
+        df_output=zero_shot_classification_fillmask_onnx(premise=input_texts,
+                                                         labels=input_lables,
+                                                         hypothesis=input_hypothesis,
+                                                         _session=session_zs_mlm,
+                                                        _tokenizer=tokenizer_zs_mlm,
+                                                         )
+        end=time.time()
+        st.write(f"Time taken for computation {(end - start) * 1000:.1f} ms")
+
+        fig = px.bar(x='Probability',
+                     y='Labels',
+                     text='Probability',
+                     data_frame=df_output,
+                     title='Zero Shot MLM Normalized Probabilities')
 
         st.plotly_chart(fig, config=_plotly_config)
     else:
         pass
 
+

config.yaml

@@ -12,3 +12,10 @@ ZEROSHOT_CLF:
   zs_onnx_mdl_name: 'model.onnx'
   zs_onnx_quant_mdl_name: 'model_quant.onnx'
 
+ZEROSHOT_MLM:
+  zs_mlm_chkpt: 'bert-base-uncased'
+  zs_mlm_mdl_dir: 'zs_mlm_dir'
+  zs_mlm_onnx_mdl_dir: 'zs_mlm_onnx_dir'
+  zs_mlm_onnx_mdl_name: 'model.onnx'
+
+

zeroshot_clf_helper.py

@@ -4,6 +4,10 @@ import os
 import subprocess
 import numpy as np
 import pandas as pd
+import transformers
+import transformers.convert_graph_to_onnx as onnx_convert
+from pathlib import Path
+import streamlit as st
 
 import yaml
 def read_yaml(file_path):
@@ -18,8 +22,24 @@ zs_onnx_mdl_dir=config['ZEROSHOT_CLF']['zs_onnx_mdl_dir']
 zs_onnx_mdl_name=config['ZEROSHOT_CLF']['zs_onnx_mdl_name']
 zs_onnx_quant_mdl_name=config['ZEROSHOT_CLF']['zs_onnx_quant_mdl_name']
 
+zs_mlm_chkpt=config['ZEROSHOT_MLM']['zs_mlm_chkpt']
+zs_mlm_mdl_dir=config['ZEROSHOT_MLM']['zs_mlm_mdl_dir']
+zs_mlm_onnx_mdl_dir=config['ZEROSHOT_MLM']['zs_mlm_onnx_mdl_dir']
+zs_mlm_onnx_mdl_name=config['ZEROSHOT_MLM']['zs_mlm_onnx_mdl_name']
+
 
 def zero_shot_classification(premise: str, labels: str, model, tokenizer):
+    """
+
+    Args:
+        premise:
+        labels:
+        model:
+        tokenizer:
+
+    Returns:
+
+    """
     try:
         labels=labels.split(',')
         labels=[l.lower() for l in labels]
@@ -49,11 +69,19 @@ def zero_shot_classification(premise: str, labels: str, model, tokenizer):
     return df
 
 ##example
-# zero_shot_classification(premise='Tiny worms and breath analyzers could screen for \disease while it’s early and treatable',
+# zero_shot_classification(premise='Tiny worms and breath analyzers could screen for disease while it’s early and treatable',
 #                          labels='science, sports, museum')
 
 
-def create_onnx_model_zs(zs_onnx_mdl_dir=zs_onnx_mdl_dir):
+def create_onnx_model_zs_nli(zs_onnx_mdl_dir=zs_onnx_mdl_dir):
+    """
+
+    Args:
+        zs_onnx_mdl_dir:
+
+    Returns:
+
+    """
 
     # create onnx model using
     if not os.path.exists(zs_onnx_mdl_dir):
@@ -61,6 +89,7 @@ def create_onnx_model_zs(zs_onnx_mdl_dir=zs_onnx_mdl_dir):
             subprocess.run(['python3', '-m', 'transformers.onnx',
                             '--model=valhalla/distilbart-mnli-12-1',
                             '--feature=sequence-classification',
+                            '--atol=1e-3',
                             zs_onnx_mdl_dir])
         except Exception as e:
             print(e)
@@ -72,7 +101,19 @@ def create_onnx_model_zs(zs_onnx_mdl_dir=zs_onnx_mdl_dir):
     else:
         pass
 
-def zero_shot_classification_onnx(premise,labels,_session,_tokenizer):
+def zero_shot_classification_nli_onnx(premise,labels,_session,_tokenizer,hypothesis="This is an example of"):
+    """
+
+    Args:
+        premise:
+        labels:
+        _session:
+        _tokenizer:
+        hypothesis:
+
+    Returns:
+
+    """
     try:
         labels=labels.split(',')
         labels=[l.lower() for l in labels]
@@ -85,7 +126,7 @@ def zero_shot_classification_onnx(premise,labels,_session,_tokenizer):
 
     for l in labels:
 
-        hypothesis= f'this is an example of {l}'
+        hypothesis= f"{hypothesis} {l}"
 
         inputs = _tokenizer(premise,hypothesis,
                              return_tensors='pt',
@@ -109,4 +150,84 @@ def zero_shot_classification_onnx(premise,labels,_session,_tokenizer):
     return df
 
 
+def create_onnx_model_zs_mlm(_model, _tokenizer,zs_mlm_onnx_mdl_dir=zs_mlm_onnx_mdl_dir):
+    """
+
+    Args:
+        _model:
+        _tokenizer:
+        zs_mlm_onnx_mdl_dir:
+
+    Returns:
+
+    """
+    if not os.path.exists(zs_mlm_onnx_mdl_dir):
+        try:
+            subprocess.run(['python3', '-m', 'transformers.onnx',
+                            f'--model={zs_mlm_chkpt}',
+                            '--feature=masked-lm',
+                            zs_mlm_onnx_mdl_dir])
+        except:
+            pass
+
+    else:
+        pass
+
+def zero_shot_classification_fillmask_onnx(premise,hypothesis,labels,_session,_tokenizer):
+    """
+
+    Args:
+        premise:
+        hypothesis:
+        labels:
+        _session:
+        _tokenizer:
+
+    Returns:
+
+    """
+    try:
+        labels=labels.split(',')
+        labels=[l.lower().rstrip().lstrip() for l in labels]
+    except:
+        raise Exception("please pass atleast 2 labels to classify")
+
+    premise=premise.lower()
+    hypothesis=hypothesis.lower()
+
+    final_input= f"{premise}.{hypothesis} [MASK]" #this can change depending on chkpt, this is for bert-base-uncased chkpt
+
+    _inputs=_tokenizer(final_input,padding=True, truncation=True,
+                      return_tensors="pt")
+
+    input_feed={
+        'input_ids': np.array(_inputs['input_ids']),
+        'token_type_ids': np.array(_inputs['token_type_ids']),
+        'attention_mask': np.array(_inputs['attention_mask'])
+    }
+
+    output=_session.run(output_names=['logits'],input_feed=dict(input_feed))[0]
+
+    mask_token_index = np.argwhere(_inputs["input_ids"] == _tokenizer.mask_token_id)[1,0]
+
+    mask_token_logits=output[0,mask_token_index,:]
+
+    #seacrh for logits of input labels
+    #encode the labels and get the label id -
+    labels_logits=[]
+    for l in labels:
+        encoded_label=_tokenizer.encode(l)[1]
+        labels_logits.append(mask_token_logits[encoded_label])
+
+    #do a softmax on the logits
+    labels_logits=np.array(labels_logits)
+    labels_logits=torch.from_numpy(labels_logits)
+    labels_logits=labels_logits.softmax(dim=0)
+
+    output= {'Labels':labels,
+            'Probability':labels_logits}
+
+    df_output = pd.DataFrame(output)
+    df_output['Probability'] = df_output['Probability'].apply(lambda x: np.round(100*x, 1))
 
+    return df_output

zs_mlm_dir/config.json

@@ -0,0 +1,26 @@
+{
+  "_name_or_path": "bert-base-uncased",
+  "architectures": [
+    "BertForMaskedLM"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "classifier_dropout": null,
+  "gradient_checkpointing": false,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 768,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "layer_norm_eps": 1e-12,
+  "max_position_embeddings": 512,
+  "model_type": "bert",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "pad_token_id": 0,
+  "position_embedding_type": "absolute",
+  "torch_dtype": "float32",
+  "transformers_version": "4.18.0",
+  "type_vocab_size": 2,
+  "use_cache": true,
+  "vocab_size": 30522
+}

zs_mlm_dir/special_tokens_map.json

@@ -0,0 +1 @@
+{"unk_token": "[UNK]", "sep_token": "[SEP]", "pad_token": "[PAD]", "cls_token": "[CLS]", "mask_token": "[MASK]"}

zs_mlm_dir/tokenizer_config.json

@@ -0,0 +1 @@
+{"do_lower_case": true, "unk_token": "[UNK]", "sep_token": "[SEP]", "pad_token": "[PAD]", "cls_token": "[CLS]", "mask_token": "[MASK]", "tokenize_chinese_chars": true, "strip_accents": null, "model_max_length": 512, "special_tokens_map_file": null, "name_or_path": "bert-base-uncased", "tokenizer_class": "BertTokenizer"}

zs_mlm_onnx_dir/model.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e55575cce3f2b3b68e82e4dbdaabff3a8a5eaaeac4703e4000b1cb717174543a
+size 531893756