Update app.py

app.py

@@ -12,7 +12,7 @@ from DecompX.src.modeling_bert import BertForSequenceClassification
 from DecompX.src.modeling_roberta import RobertaForSequenceClassification
 
 plt.style.use("ggplot")
-MODELS = ["WillHeld/roberta-base-sst2"]
+MODELS = ['TehranNLP-org/bert-base-uncased-cls-sst2', 'TehranNLP-org/bert-large-sst2', "WillHeld/roberta-base-sst2"]
 
 def plot_clf(tokens, logits, label_names, title="", file_name=None):
     print(tokens)
@@ -140,7 +140,7 @@ def run_decompx(text, model):
     elif "bert" in MODEL:
         model = BertForSequenceClassification.from_pretrained(MODEL)
     else:
-        raise Exception(f"Not implemented model: {MODEL}")
+        raise Exception(f"Not implented model: {MODEL}")
     # RUN DECOMPX
     with torch.no_grad():
         model.eval()
@@ -185,10 +185,13 @@ demo = gr.Interface(
         gr.components.Dropdown(label="Model", choices=MODELS),
     ],
     outputs=["plot", "html"],
-    examples=[["Building a translation demo with Gradio is so easy!", "WillHeld/roberta-base-sst2"]],
-    cache_examples=False,
+    examples=[
+        ["a good piece of work more often than not.", "TehranNLP-org/bert-base-uncased-cls-sst2"], 
+        ["A deep and meaningful film.", "TehranNLP-org/bert-large-sst2"],
+    ],
+    cache_examples=True,
     title="DecompX Demo",
-    description="This demo is a simplified version of the original [NLLB-Translator](https://huggingface.co/spaces/Narrativaai/NLLB-Translator) space"
+    description="This is a demo for the ACL 2023 paper [DecompX](https://github.com/mohsenfayyaz/DecompX/)"
 )
 
 demo.launch()