Update app.py

app.py

@@ -1,4 +1,6 @@
 import streamlit as st
+import pandas as pd
+import fitz  # PyMuPDF
 from transformers import pipeline
 
 # Load pre-trained model and tokenizer from Hugging Face
@@ -12,22 +14,54 @@ labels = {
 }
 
 # Streamlit app
-st.title("Text Classification")
+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")
 
+def classify_text(text):
+    result = pipe(text)[0]
+    label = labels.get(result['label'], result['label'])
+    score = result['score']
+    return label, score
+
 if st.button("Classify"):
     if input_text:
         # Perform classification
-        result = pipe(input_text)
-        
-        # Extract label and score
-        label = labels.get(result[0]['label'], result[0]['label'])
-        score = result[0]['score']
+        label, score = classify_text(input_text)
 
         st.write(f"**Predicted Class:** {label}")
         st.write(f"**Confidence:** {score:.4f}")
     else:
         st.write("Please enter some text to classify.")
+
+# File upload section
+st.write("Upload a file for classification:")
+
+uploaded_file = st.file_uploader("Choose a file", type=["csv", "pdf"])
+
+if uploaded_file is not None:
+    if uploaded_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 uploaded_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}")