Update app.py

app.py

@@ -19,11 +19,31 @@ sub = st.write("Pick the model to use for analyzing the text!")
 button = st.button("Analyze!")
 pipe = pipeline("text-classification")
 if(button):
+    if model == "aim9061/fine-tuned-toxic-tweet-dilbert":
+      get_cats(text)
     pipe = pipeline("text-classification", model)
     results = pipe(text)
     st.write(results)
-    
+def create_dict(prob, text):
+  sorted_indices = np.argsort(prob)[-2:]
+  info = {"text": text,
+         "toxicity": labels[sorted_indices[1]],
+         "percentage1": str(round(prob[sorted_indices[1]], 3)),
+         "obscenity": labels[sorted_indices[0]],
+         "percentage2": str(round(prob[sorted_indices[0]], 3))}
+  return info
+def get_cats(text):
+  tokenizer = AutoTokenizer.from_pretrained("aim9061/fine-tuned-toxic-tweet-dilbert")
+  token = tokenizer(text, return_tensors="pt")
+  
+  outputs = model(**tokens)
+  
+  prob = torch.sigmoid(output.logits).detach().numpy()[0]
+  
+  data = create_dict(prob, text)
+  st.table(data)
+  
 #TODO: DOCUMENT CODE
-#TODO: MAKE THE OUTPUT PRETTIER
+