Update app.py

app.py

@@ -4,9 +4,17 @@ from transformers import AutoTokenizer, AutoModelForSequenceClassification
 pretrained = "rohanphadke/roberta-finetuned-triplebottomline"
 tokenizer = AutoTokenizer.from_pretrained(pretrained)
 model = AutoModelForSequenceClassification.from_pretrained(pretrained)
+threshold = 0.5
+
+labels = {0: 'people', 1: 'planet', 2:'profit'}
 
 def greet(name):
     return "Hello " + name + "!!"
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
+def predict_text(text):
+    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=128)
+    probs = torch.sigmoid(model(**inputs).logits).detach().cpu().numpy()[0]
+    return {labels[i]: float(probs[i]) for i in range(len(probs)) if probs[i] >= threshold}
+
+demo = gr.Interface(fn=predict_text, inputs="text", outputs=gr.outputs.Label(num_top_classes=3))
 demo.launch()