Update app.py

app.py

@@ -1,5 +1,3 @@
-#实现功能：1忽略奇怪符号，直接删掉 2.怎么引用fine-tune的model
-
 import streamlit as st
 from transformers import pipeline
 import re
@@ -26,6 +24,20 @@ except ValueError as e:
     st.error(f"Error loading classification model: {e}")
     classifier_loaded = False
 
+# Dictionary to map numerical labels to real labels
+label_mapping = {
+    0: "Society & Culture",
+    1: "Science & Mathematics",
+    2: "Health",
+    3: "Education & Reference",
+    4: "Computers & Internet",
+    5: "Sports",
+    6: "Business & Finance",
+    7: "Entertainment & Music",
+    8: "Family & Relationships",
+    9: "Politics & Government"
+}
+
 # Streamlit app title
 st.title("Question Rephrase and Classification")
 
@@ -51,8 +63,10 @@ if st.button("Process"):
                 results = classifier(summarized_text)[0]
                 # Find the category with the highest score
                 max_score = max(results, key=lambda x: x['score'])
+                predicted_label_index = int(max_score['label'].split('_')[-1])  # Assuming labels are like "LABEL_0", "LABEL_1", etc.
+                predicted_label = label_mapping[predicted_label_index]
                 st.write("Rephrased Text:", summarized_text)
-                st.write("Category:", max_score['label'])
+                st.write("Category:", predicted_label)
                 st.write("Score:", max_score['score'])
             except Exception as e:
                 st.error(f"Error during classification: {e}")