Update app.py

app.py

@@ -4,13 +4,13 @@ from transformers import pipeline
 from stqdm import stqdm
 from simplet5 import SimpleT5
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-from transformers import BertTokenizer
+from transformers import BertTokenizer, TFBertForSequenceClassification
 from tensorflow.keras.models import load_model
 from tensorflow.nn import softmax
 import numpy as np
 from datetime import datetime
 import logging
-from constants import sub_themes_dict
+from subconst import sub_themes_dict
 
 date = datetime.now().strftime(r"%Y-%m-%d")
 model_classes = {
@@ -64,8 +64,12 @@ def classify_category():
 
 @st.cache(allow_output_mutation=True, suppress_st_warning=True)
 def classify_sub_theme():
-    tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
-    new_model = load_model("sub_theme_model")
+    tokenizer = BertTokenizer.from_pretrained(
+        "ashhadahsan/amazon-subtheme-bert-base-finetuned"
+    )
+    new_model = TFBertForSequenceClassification.from_pretrained(
+        "ashhadahsan/amazon-subtheme-bert-base-finetuned"
+    )
     return tokenizer, new_model
 
 
@@ -96,7 +100,6 @@ if st.button("Process", type="primary"):
     cancel_button3 = st.empty()
     if uploaded_file is not None:
         if uploaded_file.name.split(".")[-1] in ["xls", "xlsx"]:
-
             df = pd.read_excel(uploaded_file, engine="openpyxl")
         if uploaded_file.name.split(".")[-1] in [".csv"]:
             df = pd.read_csv(uploaded_file)
@@ -116,7 +119,6 @@ if st.button("Process", type="primary"):
                     summary = []
 
                     for x in stqdm(range(len(text))):
-
                         if cancel_button.button("Cancel", key=x):
                             del model
                             break
@@ -185,7 +187,6 @@ if st.button("Process", type="primary"):
                     model, tokenizer = load_t5()
                     summary = []
                     for x in stqdm(range(len(text))):
-
                         if cancel_button2.button("Cancel", key=x):
                             del model, tokenizer
                             break