moved code from test.py to main.py

src/main.py

@@ -1,105 +1,277 @@
-import streamlit as st
-from transformers import pipeline
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
-
 import os
+import json
+import random
+
+import streamlit as st
+from transformers import TextClassificationPipeline, pipeline
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, DistilBertTokenizerFast, DistilBertForSequenceClassification
 
 # We'll be using Torch this time around
 import torch
 import torch.nn.functional as F
 
-model_names = {
-    "emotion":[
-        "cardiffnlp/twitter-roberta-base-sentiment",
-        "finiteautomata/beto-sentiment-analysis",
-        "bhadresh-savani/distilbert-base-uncased-emotion",
-        "siebert/sentiment-roberta-large-english"
-    ],
-    "uspto":{
-        "base":"distilbert-base-uncased",
-        "abstracts":"./models/upsto_abstracts",
-        "claims":"./models/upsto_claims"
-    }
-}
-
-def label_dictionary(model_name):
-    if model_name == "cardiffnlp/twitter-roberta-base-sentiment":
-        def twitter_roberta(label):
-            if label == "LABEL_0":
-                return "Negative"
-            elif label == "LABEL_2":
-                return "Positive"
-            else:
-                return "Neutral"
-        return twitter_roberta
-    return lambda x: x
-
-@st.cache(allow_output_mutation=True)
-def load_model(model_name):
-    model = AutoModelForSequenceClassification.from_pretrained(model_name)
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    classifier = pipeline(task="sentiment-analysis", model=model, tokenizer=tokenizer)
-    parser = label_dictionary(model_name)
-    return model, tokenizer, classifier, parser
-
-# We first initialize a state. The state will include the following:
-# 1) the name of the model (default: cardiffnlp/twitter-roberta-base-sentiment)
-# 2) the model itself, and 
-# 3) the parser for the outputs, in case we actually need to parse the output to something more sensible
-if "model" not in st.session_state:
-    st.session_state.model_name = "cardiffnlp/twitter-roberta-base-sentiment"
-    model, tokenizer, classifier, label_parser = load_model("cardiffnlp/twitter-roberta-base-sentiment")
-    st.session_state.model = model
-    st.session_state.tokenizer = tokenizer
-    st.session_state.classifier = classifier
-    st.session_state.label_parser = label_parser
-    st.session_state.panel = "emotion"
-
-def model_change():
-    model, tokenizer, classifier, label_parser = load_model(st.session_state.model_name)
-    st.session_state.model = model
-    st.session_state.tokenizer = tokenizer
-    st.session_state.classifier = classifier
-    st.session_state.label_parser = label_parser
+emotion_model_names = (
+    "cardiffnlp/twitter-roberta-base-sentiment",
+    "finiteautomata/beto-sentiment-analysis",
+    "bhadresh-savani/distilbert-base-uncased-emotion",
+    "siebert/sentiment-roberta-large-english"
+)
+
+class ModelImplementation(object):
+    def __init__(
+            self, 
+            transformer_model_name, 
+            model_transformer, 
+            tokenizer_model_name,
+            tokenizer_func, 
+            pipeline_func, 
+            parser_func,
+            classifier_args={},
+            placeholders=[""]
+    ):
+        self.transformer_model_name = transformer_model_name
+        self.tokenizer_model_name = tokenizer_model_name
+        self.placeholders = placeholders
+
+        self.model = model_transformer.from_pretrained(self.transformer_model_name)
+        self.tokenizer = tokenizer_func.from_pretrained(self.tokenizer_model_name)
+        self.classifier = pipeline_func(model=self.model, tokenizer=self.tokenizer, padding=True, truncation=True, **classifier_args)
+        self.parser = parser_func
+
+        self.history = []
+    
+    def predict(self, val):
+        result = self.classifier(val)
+        return self.parser(self, result)
+
+def ParseEmotionOutput(self, result):
+    label = result[0]['label']
+    score = result[0]['score']
+    if self.transformer_model_name == "cardiffnlp/twitter-roberta-base-sentiment":
+        if label == "LABEL_0":
+            label = "Negative"
+        elif label == "LABEL_2":
+            label = "Positive"
+        else:
+            label = "Neutral"
+    return label, score
+
+def ParsePatentOutput(self, result):
+    return result
+
+def emotion_model_change():
+    st.session_state.emotion_model = ModelImplementation(
+        st.session_state.emotion_model_name,
+        AutoModelForSequenceClassification,
+        st.session_state.emotion_model_name,
+        AutoTokenizer,
+        pipeline,
+        ParseEmotionOutput,
+        classifier_args={ "task" : "sentiment-analysis" },
+        placeholders=["@AmericanAir just landed - 3hours Late Flight - and now we need to wait TWENTY MORE MINUTES for a gate! I have patience but none for incompetence."]
+    )
+
+if "page" not in st.session_state:
+    st.session_state.page = "home"
+
+if "emotion_model_name" not in st.session_state:
+    st.session_state.emotion_model_name = "cardiffnlp/twitter-roberta-base-sentiment"
+    emotion_model_change()
+
+if "patent_data" not in st.session_state:
+    f = open('./data/val.json')
+    valData = json.load(f)
+    f.close()
+
+    patent_data = {}
+    for num, label, abstract, claim in zip(valData["patent_numbers"],valData["labels"], valData["abstracts"], valData["claims"]):
+        patent_data[num] = {"patent_number":num,"label":label,"abstract":abstract,"claim":claim}
+
+    st.session_state.patent_data = patent_data
+    st.session_state.patent_num = list(patent_data.keys())[0]
+    st.session_state.weight = 0.5
+    st.session_state.patent_abstract_model = ModelImplementation(
+        './models/uspto_abstracts',
+        DistilBertForSequenceClassification,
+        'distilbert-base-uncased',
+        DistilBertTokenizerFast,
+        TextClassificationPipeline,
+        ParsePatentOutput,
+        classifier_args={"return_all_scores":True},
+    )
+    print("Patent abstracts model initialized")
+    st.session_state.patent_claim_model = ModelImplementation(
+        './models/uspto_claims',
+        DistilBertForSequenceClassification,
+        'distilbert-base-uncased',
+        DistilBertTokenizerFast,
+        TextClassificationPipeline,
+        ParsePatentOutput,
+        classifier_args={"return_all_scores":True},
+    )
+    print("Patent claims model initialized")
 
 # Title
-st.title("CSGY-6613 Sentiment Analysis")
+st.title("CSGY-6613 Project")
 # Subtitle
-st.markdown("### Ryan Kim (rk2546)")
-st.markdown("")
-
-model_option = st.selectbox(
-    "What sentiment analysis model do you want to use?",
-    (
-        "cardiffnlp/twitter-roberta-base-sentiment",
-        "finiteautomata/beto-sentiment-analysis",
-        "bhadresh-savani/distilbert-base-uncased-emotion",
-        "siebert/sentiment-roberta-large-english"
-    ),
-    on_change=model_change,
-    key="model_name"
-)
-placeholder="@AmericanAir just landed - 3hours Late Flight - and now we need to wait TWENTY MORE MINUTES for a gate! I have patience but none for incompetence."
-form = st.form(key='sentiment-analysis-form')
-text_input = form.text_area("Enter some text for sentiment analysis! If you just want to test it out without entering anything, just press the \"Submit\" button and the model will look at the placeholder.", placeholder=placeholder)
-submit = form.form_submit_button('Submit')
-
-if submit:
-    if text_input is None or len(text_input.strip()) == 0:
-        to_eval = placeholder
+st.markdown("_**Ryan Kim (rk2546)**_")
+st.markdown("---")
+
+def PageToHome():
+    st.session_state.page = "home"
+def PageToEmotion():
+    st.session_state.page = "emotion"
+def PageToPatent():
+    st.session_state.page = "patent"
+
+with st.sidebar:
+    st.subheader("Toolbox")
+    home_selected = st.button("Home", on_click=PageToHome)
+    emotion_selected = st.button(
+        "Emotion Analysis [Milestone #2]", 
+        on_click=PageToEmotion
+    )
+    patent_selected = st.button(
+        "Patent Prediction [Milestone #3]", 
+        on_click=PageToPatent
+    )
+
+if st.session_state.page == "emotion":
+    st.subheader("Sentiment Analysis")
+    if "emotion_model" not in st.session_state:
+        st.write("Loading model...")
     else:
-        to_eval = text_input.strip()
-    st.write("You entered:")
-    st.markdown("> {}".format(to_eval))
-    st.write("Using the NLP model:")
-    st.markdown("> {}".format(st.session_state.model_name))
-    result = st.session_state.classifier(to_eval)  
-    label = result[0]['label']
-    score = result[0]['score']
+        model_option = st.selectbox(
+            "What sentiment analysis model do you want to use? NOTE: Lag may occur when loading a new model!",
+            emotion_model_names,
+            on_change=emotion_model_change,
+            key="emotion_model_name"
+        )
+        form = st.form(key='sentiment-analysis-form')
+        text_input = form.text_area(
+            "Enter some text for sentiment analysis! If you just want to test it out without entering anything, just press the \"Submit\" button and the model will look at the placeholder.", 
+            placeholder=st.session_state.emotion_model.placeholders[0]
+        )
+        submit = form.form_submit_button('Submit')
+        if submit:
+            if text_input is None or len(text_input.strip()) == 0:
+                to_eval = st.session_state.emotion_model.placeholders[0]
+            else:
+                to_eval = text_input.strip()
+            st.write("You entered:")
+            st.markdown("> {}".format(to_eval))
+            st.write("Using the NLP model:")
+            st.markdown("> {}".format(st.session_state.emotion_model_name))
+            label, score = st.session_state.emotion_model.predict(to_eval)  
+            st.markdown("#### Result:")
+            st.markdown("**{}**: {}".format(label,score))
+    
+elif st.session_state.page == "patent":
+    st.subheader("USPTO Patent Evaluation")
+    st.markdown("Below are two inputs - one for an **ABSTRACT** and another for a list of **CLAIMS**. Enter both and select the \"Submit\" button to evaluate the patenteability of your idea.")
+
+    patent_select_list = list(st.session_state.patent_data.keys())
+    patent_index_option = st.selectbox(
+        "Want to pre-populate with an existing patent? Select the index number of  below.",
+        patent_select_list,
+        key="patent_num",
+    )
+
+    print(patent_index_option)
+
+    if "patent_abstract_model" not in st.session_state or "patent_claim_model" not in st.session_state:
+        st.write("Loading models...")
+    else:
+        with st.form(key='patent-form'):
+            col1, col2 = st.columns(2)
+            with col1:
+                abstract_input = st.text_area(
+                    "Enter the abstract of the patent below", 
+                    placeholder=st.session_state.patent_data[st.session_state.patent_num]["abstract"],
+                    height=400
+                )
+            with col2:
+                claim_input = st.text_area(
+                    "Enter the claims of the patent below", 
+                    placeholder=st.session_state.patent_data[st.session_state.patent_num]["claim"],
+                    height=400
+                )
+            weight_val = st.slider(
+                "How much do the abstract and claims weight when aggregating a total softmax score?",
+                min_value=-1.0,
+                max_value=1.0,
+                value=0.5,
+            )
+            submit = st.form_submit_button('Submit')
+
+            if submit:
+                
+                is_custom = False
+                if abstract_input is None or len(abstract_input.strip()) == 0:
+                    abstract_to_eval = st.session_state.patent_data[st.session_state.patent_num]["abstract"].strip()
+                else:
+                    abstract_to_eval = abstract_input.strip()
+                    is_custom = True
+
+                if claim_input is None or len(claim_input.strip()) == 0:
+                    claim_to_eval = st.session_state.patent_data[st.session_state.patent_num]["claim"].strip()
+                else:
+                    claim_to_eval = claim_input.strip()
+                    is_custom = True
+
+                #tokenized_claim = st.session_state.patent_claim_model.tokenizer.encode(claim_to_eval, padding=True, truncation=True, max_length=512, add_special_tokens = True)
+                #untokenized_claim = st.session_state.patent_claim_model.tokenizer.decode(tokenized_claim)
+                #claim_to_eval2 = untokenized_claim.replace("[CLS]","")
+                #claim_to_eval2 = claim_to_eval2.replace("[SEP]","")
+                #print(claim_to_eval2)
+
+                abstract_response = st.session_state.patent_abstract_model.predict(abstract_to_eval)
+                claim_response = st.session_state.patent_claim_model.predict(claim_to_eval)
+                print(abstract_response[0])
+                print(claim_response[0])
+                print(weight_val)
+
+                claim_weight = (1+weight_val)/2
+                abstract_weight = 1-claim_weight
+                aggregate_score = [
+                    {'label':'REJECTED','score':abstract_response[0][0]['score']*abstract_weight + claim_response[0][0]['score']*claim_weight},
+                    {'label':'ACCEPTED','score':abstract_response[0][1]['score']*abstract_weight + claim_response[0][1]['score']*claim_weight}
+                ]
+                aggregate_score_sorted = sorted(aggregate_score, key=lambda d: d['score'], reverse=True) 
+                print(aggregate_score_sorted)
+                print(f'Original Rating: {st.session_state.patent_data[st.session_state.patent_num]["label"]}')
+
+                st.markdown("---")
+                answerCol1, answerCol2 = st.columns(2)
+                with answerCol1:
+                    st.markdown("### Abstract Ratings")
+                    st.markdown("""
+                        > **Reject**: {}  
+                        > **Accept**: {}
+                    """.format(abstract_response[0][0]["score"], abstract_response[0][1]["score"]))
+                with answerCol2:
+                    st.markdown("### Claims Ratings")
+                    st.markdown("""
+                        > **Reject**: {}  
+                        > **Accept**: {}
+                    """.format(claim_response[0][0]["score"], claim_response[0][1]["score"]))
+                
+                st.markdown(f'### Final Rating: **{aggregate_score_sorted[0]["label"]}**')
+                st.markdown("""
+                    > **Reject**: {}  
+                    > **Accept**: {}
+                """.format(aggregate_score[0]['score'], aggregate_score[1]['score']))
+                
+                #if not is_custom:
+                #    st.markdown('**Original Score:**')
+                #    st.markdown(st.session_state.patent_data[st.session_state.patent_num]["label"])
+                
+
+            
+
 
-    label = st.session_state.label_parser(label)
+else:
+    st.write("To get started, access the sidebar on the left (click the arrow in the top-left corner of the screen) and select a tool.")
 
-    st.markdown("#### Result:")
-    st.markdown("**{}**: {}".format(label,score))
-    st.write("")
-    st.write("")
+st.write("")
+st.write("")