Update app.py

app.py

@@ -70,12 +70,14 @@ def main():
             predict_proba = round(torch.softmax(logits, 1).cpu().squeeze().tolist()[prediction],4)
             return prediction, predict_proba
 
-    def pred_str(prediction):
+    def pred_str(prediction, option):
     # Map the predicted class to string output
         if prediction == 0:
             return "Human-made 🤷‍♂️🤷‍♀️"
+        elif prediction == 1 and "Revised" in option:
+            return "Revised with AI 🦾"
         else:
-            return "Generated/Revised with AI 🦾"
+            return "Generated with AI 🦾"
     
     @st.cache(allow_output_mutation=True, suppress_st_warning=True)
     def load_tokenizer(option):
@@ -157,7 +159,7 @@ def main():
 
             # Store the result in session state
             st.session_state["color_pred"] = "blue" if prediction == 0 else "red"
-            prediction = pred_str(prediction)
+            prediction = pred_str(prediction, option)
             st.session_state["prediction"] = prediction
             st.session_state["predict_proba"] = predict_proba
             st.session_state["text"] = text