Update app.py

app.py

@@ -5,27 +5,19 @@ import docx
 from transformers import BertTokenizer, BertForSequenceClassification
 import torch
 from io import StringIO
-import openpyxl
 
-# Load pre-trained model and tokenizer from Hugging Face
+# Load pre-trained model and tokenizer
 model_name = "google/bert-base-uncased"
 tokenizer = BertTokenizer.from_pretrained(model_name)
 model = BertForSequenceClassification.from_pretrained(model_name, num_labels=2)
 
-# Custom labels for your classification task
+# Custom labels for classification
 labels = {
-    "LABEL_0": "Negative",
-    "LABEL_1": "Positive"
+    0: "Negative",
+    1: "Positive"
 }
 
-# Streamlit app
-st.title("BERT Text Classification")
-
-st.write("This app uses a pre-trained BERT model to classify text into positive or negative sentiment.")
-
-# Input text area
-input_text = st.text_area("Enter text to classify")
-
+# Function to classify text
 def classify_text(text):
     inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
     outputs = model(**inputs)
@@ -35,7 +27,7 @@ def classify_text(text):
     confidence_score = confidence_score.item()
     predicted_class = predicted_class.item()
 
-    label = labels[f"LABEL_{predicted_class}"]
+    label = labels[predicted_class]
     
     # Adjust classification based on score
     if confidence_score < 0.75:
@@ -43,73 +35,37 @@ def classify_text(text):
     
     return label, confidence_score
 
-if st.button("Classify"):
-    if input_text:
-        # Perform classification
-        label, score = classify_text(input_text)
+# Streamlit app
+st.title("BERT Text Classification")
+st.write("This app classifies text using a pre-trained BERT model.")
+
+# Input text area
+input_text = st.text_area("Enter text to classify")
 
-        st.write(f"**Predicted Class:** {label}")
-        st.write(f"**Confidence:** {score:.4f}")
-    else:
-        st.write("Please enter some text to classify.")
+# Classification and display
+if st.button("Classify") and input_text:
+    label, score = classify_text(input_text)
+    st.write(f"Predicted Class: {label}")
+    st.write(f"Confidence: {score:.4f}")
 
 # File upload section
 st.write("Upload a file for classification:")
-
 uploaded_file = st.file_uploader("Choose a file", type=["csv", "pdf", "txt", "doc", "docx", "xlsx"])
 
 if uploaded_file is not None:
     file_type = uploaded_file.type
     if file_type == "text/csv":
-        # Process CSV file
         df = pd.read_csv(uploaded_file)
-        if 'text' not in df.columns:
-            st.write("The CSV file must contain a 'text' column.")
-        else:
-            df['Prediction'] = df['text'].apply(lambda x: classify_text(x)[0])
-            df['Confidence'] = df['text'].apply(lambda x: classify_text(x)[1])
-            st.write(df)
-
     elif file_type == "application/pdf":
-        # Process PDF file
         with fitz.open(stream=uploaded_file.read(), filetype="pdf") as doc:
-            text = ""
-            for page in doc:
-                text += page.get_text()
-
-            # Perform classification
-            label, score = classify_text(text)
-
-            st.write(f"**Predicted Class for PDF:** {label}")
-            st.write(f"**Confidence:** {score:.4f}")
-
+            text = "".join(page.get_text() for page in doc)
     elif file_type == "text/plain":
-        # Process TXT file
-        text = StringIO(uploaded_file.getvalue().decode("utf-8")).read()
-
-        # Perform classification
-        label, score = classify_text(text)
-
-        st.write(f"**Predicted Class for TXT:** {label}")
-        st.write(f"**Confidence:** {score:.4f}")
-
+        text = uploaded_file.getvalue().decode("utf-8")
     elif file_type in ["application/vnd.openxmlformats-officedocument.wordprocessingml.document", "application/msword"]:
-        # Process DOCX or DOC file
         doc = docx.Document(uploaded_file)
-        text = "\n".join([para.text for para in doc.paragraphs])
+        text = "\n".join(para.text for para in doc.paragraphs)
 
-        # Perform classification
+    if 'text' in locals():
         label, score = classify_text(text)
-
-        st.write(f"**Predicted Class for DOC/DOCX:** {label}")
-        st.write(f"**Confidence:** {score:.4f}")
-
-    elif file_type == "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet":
-        # Process XLSX file
-        df = pd.read_excel(uploaded_file)
-        if 'text' not in df.columns:
-            st.write("The XLSX file must contain a 'text' column.")
-        else:
-            df['Prediction'] = df['text'].apply(lambda x: classify_text(x)[0])
-            df['Confidence'] = df['text'].apply(lambda x: classify_text(x)[1])
-            st.write(df)
+        st.write(f"Predicted Class for {file_type}: {label}")
+        st.write(f"Confidence: {score:.4f}")