Ak28Akhil/TextClassification

1import streamlit as st2import pandas as pd3import fitz  # PyMuPDF4from transformers import pipeline5 6# Load pre-trained model and tokenizer from Hugging Face7model_name = "google-bert/bert-base-uncased"8pipe = pipeline("text-classification", model=model_name)9 10# Custom labels for your classification task11labels = {12    "LABEL_0": "Negative",13    "LABEL_1": "Positive"14}15 16# Streamlit app17st.title("BERT Text Classification")18st.write("This app uses a pre-trained BERT model to classify text into positive or negative sentiment.")19 20# Input text area21input_text = st.text_area("Enter text to classify")22 23def classify_text(text):24    result = pipe(text)[0]25    label = labels.get(result['label'], result['label'])26    score = result['score']27    28    # Adjust classification based on score29    if score < 0.75:30        label = "Negative"31    32    return label, score33 34if st.button("Classify"):35    if input_text:36        # Perform classification37        label, score = classify_text(input_text)38        st.write(f"**Predicted Class:** {label}")39        st.write(f"**Confidence:** {score:.4f}")40    else:41        st.write("Please enter some text to classify.")42 43# File upload section44st.write("Upload a file for classification:")45uploaded_file = st.file_uploader("Choose a file", type=["csv", "pdf"])46 47if uploaded_file is not None:48    try:49        if uploaded_file.type == "text/csv":50            # Process CSV file51            df = pd.read_csv(uploaded_file, encoding='utf-8')52            if 'text' not in df.columns:53                st.write("The CSV file must contain a 'text' column.")54            else:55                df['Prediction'] = df['text'].apply(lambda x: classify_text(x)[0])56                df['Confidence'] = df['text'].apply(lambda x: classify_text(x)[1])57                st.write(df)58 59        elif uploaded_file.type == "application/pdf":60            # Process PDF file61            with fitz.open(stream=uploaded_file.read(), filetype="pdf") as doc:62                text = ""63                for page in doc:64                    text += page.get_text()65 66                # Perform classification67                label, score = classify_text(text)68                st.write(f"**Predicted Class for PDF:** {label}")69                st.write(f"**Confidence:** {score:.4f}")70    except Exception as e:71        st.error(f"Error: {e}")72