Update app.py

app.py

@@ -37,9 +37,10 @@ group_dict = {
     28: 'Small-scale farmers',
     29: 'Other'}
 
-def predict(text):
-    preds = model([text])[0].item()
-    return group_dict[preds]
+#def predict(text):
+ #   preds = model([text])[0].item()
+  #  return group_dict[preds]
+
 
 # App 
 st.title("Identify references to vulnerable groups.")
@@ -47,6 +48,16 @@ st.write("This app allows you to identify whether a text contains any references
 
 #col1, col2 = st.columns(2)
 
+# Create text input box
 input_text = st.text_area('Please enter your text here')
 
-st.text(predict(input_text))
+# Make predictions
+preds = model(input_text)
+
+# Select lab
+def get_label(prediction_tensor):
+    key = prediction_tensor.index(1)
+    return group_dict[key]
+    
+
+st.text(get_label(preds))