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()