Update app.py

app.py

@@ -1,65 +1,68 @@
 import streamlit as st
 import pandas as pd
 import fitz  # PyMuPDF
-import docx
-from transformers import BertTokenizer, BertForSequenceClassification
-import torch
-from io import StringIO
+from transformers import pipeline
 
-# 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)
+# Load pre-trained model and tokenizer from Hugging Face
+model_name = "google-bert/bert-base-uncased"
+pipe = pipeline("text-classification", model=model_name)
 
-# Custom labels for classification
+# Custom labels for your classification task
 labels = {
-    0: "Negative",
-    1: "Positive"
+    "LABEL_0": "Negative",
+    "LABEL_1": "Positive"
 }
 
-# Function to classify text
-def classify_text(text):
-    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
-    outputs = model(**inputs)
-    logits = outputs.logits
-    probabilities = torch.softmax(logits, dim=-1)
-    confidence_score, predicted_class = torch.max(probabilities, dim=-1)
-    confidence_score = confidence_score.item()
-    predicted_class = predicted_class.item()
-
-    label = labels[predicted_class]
-    
 # Streamlit app
 st.title("BERT Text Classification")
-st.write("This app classifies text using a pre-trained BERT model.")
+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")
 
-# 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}")
+def classify_text(text):
+    result = pipe(text)[0]
+    label = labels.get(result['label'], result['label'])
+    score = result['score']
+    
+    # Adjust classification based on score
+    if score < 0.75:
+        label = "Negative"
+    
+    return label, score
+
+if st.button("Classify"):
+    if input_text:
+        # Perform classification
+        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", "txt", "doc", "docx", "xlsx"])
+uploaded_file = st.file_uploader("Choose a file", type=["csv", "pdf"])
 
 if uploaded_file is not None:
-    file_type = uploaded_file.type
-    if file_type == "text/csv":
+    if uploaded_file.type == "text/csv":
+        # Process CSV file
         df = pd.read_csv(uploaded_file)
-    elif file_type == "application/pdf":
+        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 = "".join(page.get_text() for page in doc)
-    elif file_type == "text/plain":
-        text = uploaded_file.getvalue().decode("utf-8")
-    elif file_type in ["application/vnd.openxmlformats-officedocument.wordprocessingml.document", "application/msword"]:
-        doc = docx.Document(uploaded_file)
-        text = "\n".join(para.text for para in doc.paragraphs)
+            text = ""
+            for page in doc:
+                text += page.get_text()
 
-    if 'text' in locals():
-        label, score = classify_text(text)
-        st.write(f"Predicted Class for {file_type}: {label}")
-        st.write(f"Confidence: {score:.4f}")
+            # Perform classification
+            label, score = classify_text(text)
+            st.write(f"**Predicted Class for PDF:** {label}")
+            st.write(f"**Confidence:** {score:.4f}")