Update app.py

app.py

@@ -56,8 +56,10 @@ def predict_text(params):
         pred = bert_model.predict([input_ids, token_type_ids, attention_mask])
         logits = pred.logits
         pred_label = tf.argmax(logits, axis=1).numpy()[0]
+        if not pred_label:
+            predict_label = 7
         
-        label = {0: 'BUSINESS', 1: 'COMEDY', 2: 'CRIME', 3: 'FOOD & DRINK', 4: 'POLITICS', 5: 'SPORTS', 6: 'TRAVEL'}
+        label = {0: 'BUSINESS', 1: 'COMEDY', 2: 'CRIME', 3: 'FOOD & DRINK', 4: 'POLITICS', 5: 'SPORTS', 6: 'TRAVEL', 7: 'None'}
         result = {'text':text, 'label':[label[pred_label]], "qcUser" : None,"normalfileID":file_id}
         solutions.append(result)