Rob Caamano commited on
Commit
ebebffc
1 Parent(s): 9cda51f

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -2
app.py CHANGED
@@ -24,7 +24,6 @@ model_name = ""
24
 
25
  model_mapping = {
26
  "Toxicity": "RobCaamano/toxicity",
27
- "Toxicity 2": "RobCaamano/toxicity_distilbert",
28
  "DistilBERT Base Uncased (SST-2)": "distilbert-base-uncased-finetuned-sst-2-english",
29
  }
30
 
@@ -47,7 +46,7 @@ input = tokenizer(text, return_tensors="tf")
47
  if submit:
48
  results = dict(d.values() for d in clf(text)[0])
49
 
50
- if model_name in ["RobCaamano/toxicity", "RobCaamano/toxicity_distilbert"]:
51
  classes = {k: results[k] for k in results.keys() if not k == "toxic"}
52
 
53
  max_class = max(classes, key=classes.get)
 
24
 
25
  model_mapping = {
26
  "Toxicity": "RobCaamano/toxicity",
 
27
  "DistilBERT Base Uncased (SST-2)": "distilbert-base-uncased-finetuned-sst-2-english",
28
  }
29
 
 
46
  if submit:
47
  results = dict(d.values() for d in clf(text)[0])
48
 
49
+ if model_name == "RobCaamano/toxicity":
50
  classes = {k: results[k] for k in results.keys() if not k == "toxic"}
51
 
52
  max_class = max(classes, key=classes.get)