Rob Caamano commited on
Commit
b081db9
1 Parent(s): 94c9a58

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -3
app.py CHANGED
@@ -14,7 +14,6 @@ text = st.text_area("Input Text", demo, height=250)
14
  model_options = {
15
  "DistilBERT Base Uncased (SST-2)": "distilbert-base-uncased-finetuned-sst-2-english",
16
  "Fine-tuned Toxicity Model": "RobCaamano/toxicity",
17
- "Fine-tuned Toxicity Model - Optimized": "RobCaamano/toxicity_optimized",
18
  }
19
  selected_model = st.selectbox("Select Model", options=list(model_options.keys()))
20
 
@@ -23,7 +22,7 @@ mod_name = model_options[selected_model]
23
  tokenizer = AutoTokenizer.from_pretrained(mod_name)
24
  model = AutoModelForSequenceClassification.from_pretrained(mod_name)
25
 
26
- if selected_model in ["Fine-tuned Toxicity Model", "Fine-tuned Toxicity Model - Optimized"]:
27
  toxicity_classes = ["toxic", "severe_toxic", "obscene", "threat", "insult", "identity_hate"]
28
  model.config.id2label = {i: toxicity_classes[i] for i in range(model.config.num_labels)}
29
 
@@ -39,7 +38,7 @@ if st.button("Submit", type="primary"):
39
 
40
  tweet_portion = text[:50] + "..." if len(text) > 50 else text
41
 
42
- if selected_model in ["Fine-tuned Toxicity Model", "Model 3.0"]:
43
  column_name = "Toxicity Class"
44
  else:
45
  column_name = "Prediction"
 
14
  model_options = {
15
  "DistilBERT Base Uncased (SST-2)": "distilbert-base-uncased-finetuned-sst-2-english",
16
  "Fine-tuned Toxicity Model": "RobCaamano/toxicity",
 
17
  }
18
  selected_model = st.selectbox("Select Model", options=list(model_options.keys()))
19
 
 
22
  tokenizer = AutoTokenizer.from_pretrained(mod_name)
23
  model = AutoModelForSequenceClassification.from_pretrained(mod_name)
24
 
25
+ if selected_model in ["Fine-tuned Toxicity Model"]:
26
  toxicity_classes = ["toxic", "severe_toxic", "obscene", "threat", "insult", "identity_hate"]
27
  model.config.id2label = {i: toxicity_classes[i] for i in range(model.config.num_labels)}
28
 
 
38
 
39
  tweet_portion = text[:50] + "..." if len(text) > 50 else text
40
 
41
+ if selected_model in ["Fine-tuned Toxicity Model"]:
42
  column_name = "Toxicity Class"
43
  else:
44
  column_name = "Prediction"