app.py

import json
import os
from typing import Tuple, Optional
import pandas as pd
import gradio as gr
import nltk
from tqdm import tqdm
from tempfile import NamedTemporaryFile
from datetime import datetime
import pytz
import smtplib
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from email.mime.base import MIMEBase
from email import encoders
from huggingface_hub import HfApi, upload_file
from dotenv import load_dotenv
load_dotenv()
api = HfApi(token=os.environ["HF_TOKEN"])

nltk.download('punkt_tab')

from transformers import AutoModelForSequenceClassification, AutoTokenizer
from transformers.pipelines import pipeline
import torch
import numpy as np
from scipy.special import softmax

# Load sentiment model and tokenizer
sentiment_model_name = "cardiffnlp/twitter-roberta-base-sentiment"
tokenizer = AutoTokenizer.from_pretrained(sentiment_model_name)
sentiment_model = AutoModelForSequenceClassification.from_pretrained(sentiment_model_name)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
sentiment_model.to(device)

# Model paths
model_path = 'Herb-Lab/LLM_housing_livability'
HF_REPO_ID = "Herb-Lab/LLM_housing_livability_space_data"

# Load models and tokenizer
classifier = pipeline("zero-shot-classification", model=model_path, device=device)

# Label set
label_to_int = ['Indoor Air Quality', 'Thermal', 'Acoustic', 'Visual']

def paragraph_to_list(paragraph : str):
    return nltk.sent_tokenize(paragraph)

def sentiment_analysis(input_text: str):
    encoded_inputs = tokenizer(input_text, return_tensors='pt', padding=True, truncation=True, max_length=512).to(device)
    with torch.no_grad():
        outputs = sentiment_model(**encoded_inputs)
    score = outputs[0].detach().cpu().numpy()
    score_softmax = softmax(score, axis=1)
    sentiment_score = float(score_softmax[0, 2] - score_softmax[0, 0])
    return sentiment_score

def upload_to_hf_dataset(file_path: str, hf_path: str):
    out = api.upload_file(
        path_or_fileobj=file_path,
        path_in_repo=hf_path,
        repo_id=HF_REPO_ID,
        repo_type="dataset"
    )
    return out

def send_email_with_attachment(recipient, file_path):
    msg = MIMEMultipart()
    msg["Subject"] = "Your labeled housing livability data is ready!"
    msg["From"] = os.environ["SMTP_FROM"]
    msg["To"] = recipient

    # Email body
    body = "Your CSV file has been processed. Please find it attached."
    msg.attach(MIMEText(body, "plain"))

    # Attach the file
    with open(file_path, "rb") as f:
        part = MIMEBase("application", "octet-stream")
        part.set_payload(f.read())
        encoders.encode_base64(part)
        part.add_header("Content-Disposition", f"attachment; filename={os.path.basename(file_path)}")
        msg.attach(part)

    # Send email
    with smtplib.SMTP(os.environ["SMTP_SERVER"], int(os.environ["SMTP_PORT"])) as smtp:
        smtp.starttls()
        smtp.login(os.environ["SMTP_USER"], os.environ["SMTP_PASS"])
        smtp.send_message(msg)

def classify_text(premise : str, sent : bool = False) -> Tuple[Optional[dict], Optional[pd.DataFrame]]:
    premise_sent_list = paragraph_to_list(premise)
    sentence_level_results = []

    # Post level classification 
    post_result = classifier(premise, label_to_int, multi_label=True)
    best_post_score = post_result['scores'][0]
    if best_post_score < 0.6: return None, None
    
    # Sentence level classification
    for _premise_sent in premise_sent_list:
        # Use the classifier pipeline for zero-shot classification
        result = classifier(_premise_sent, label_to_int, multi_label=True)
        
        # Extract the most probable label
        best_label = result['labels'][0]
        best_score = result['scores'][0]
        
        if sent: # Sentence level sentiment analysis
            sentiment_score = sentiment_analysis(_premise_sent)
            if best_score >= 0.6:
                sentence_level_results.append({
                    'post_body_sent': _premise_sent,
                    'sent_label': best_label,
                    'sent_probability': best_score,
                    'sentiment_score': sentiment_score
                })
            else:
                sentence_level_results.append({
                    'post_body_sent': _premise_sent,
                    'sent_label': 'None',
                    'sent_probability': -1,
                    'sentiment_score': sentiment_score
                })
        else: # No sentence level sentiment analysis
            if best_score >= 0.6:
                sentence_level_results.append({
                    'post_body_sent': _premise_sent,
                    'sent_label': best_label,
                    'sent_probability': best_score
                })
            else:
                sentence_level_results.append({
                    'post_body_sent': _premise_sent,
                    'sent_label': 'None',
                    'sent_probability': -1
                })
    if len(sentence_level_results) == 0: return None, None
    
    # Post level sentiment analysis 
    sentiment_score_post = None
    if sent:
        sentiment_score_post = sentiment_analysis(premise)
    
    sentence_level_results_df = pd.DataFrame(sentence_level_results)

    # Structure the output
    post_result_structured = {label: prob for label, prob in zip(post_result['labels'], post_result['scores']) if prob >= 0.6}
    if sent and sentiment_score_post is not None:
        post_result_structured['sentiment_score'] = sentiment_score_post
    return post_result_structured, sentence_level_results_df

def classify_text_for_display(premise: str) -> Tuple[Optional[dict], Optional[pd.DataFrame]]:
    post_result, sentence_results = classify_text(premise)
    if post_result is None:
        return {'Not related to housing livability': 1.0}, None
    if sentence_results is not None:
        sentence_results = sentence_results.rename(columns={
            'post_body_sent': 'Sentence of the post',
            'sent_label': 'Livability classification',
            'sent_probability': 'Confidence of classification',
        })
        sentence_results['Confidence of classification'] = sentence_results['Confidence of classification'].apply(
            lambda score: 'None' if score < 0 else f'{score:.2%}'
        )
    return post_result, sentence_results

def format_probability(probability: float) -> str:
    return f'{float(probability):.2%}'

def format_sentence_level_results(sentence_level_results: pd.DataFrame) -> str:
    pretty_results = []
    for result in sentence_level_results.to_dict(orient='records'):
        pretty_results.append({
            'Sentence of the post': result['post_body_sent'],
            'Livability classification': result['sent_label'],
            'Confidence of classification': 'None' if result['sent_probability'] < 0 else format_probability(result['sent_probability']),
            'Sentence sentiment score': round(float(result['sentiment_score']), 3),
        })
    return json.dumps(pretty_results, ensure_ascii=False)
    
def process_csv(file, user_consent=False, email_input=None):
    # Load CSV file
    try:
        df = pd.read_csv(file.name)
    except Exception as e:
        yield "Error reading the file. Please ensure it is a valid CSV file.", None, None
        return

    # Precheck: Ensure there's only one column and rename it
    if len(df.columns) != 1:
        yield "Error: The CSV file must contain exactly one column.", None, None
        return
    
    # Rename the first column to 'post_body'
    df.columns = ['post_body']

    # Save the uploaded CSV if user consents
    if user_consent:
        with NamedTemporaryFile(delete=False, suffix=".csv") as original_temp_file:
            df.to_csv(original_temp_file.name, index=False)
            original_temp_path = original_temp_file.name

        try:
            # Get current time in EST
            est = pytz.timezone("US/Eastern")
            timestamp = datetime.now(est).strftime("%Y-%m-%d_%H-%M-%S")
            if email_input is not None:
                filename = f"input_data_{email_input}_{timestamp}.csv"
            else:
                filename = f"input_data_{timestamp}.csv"
            upload_result = upload_to_hf_dataset(original_temp_path, f"bulk_uploads/{filename}")
            print(f"Bulk upload saved to Hugging Face: {upload_result}")
        except Exception as e:
            print("Error saving original file to hf:", e) #TODO

    # Limit to the first 100 rows
    df = df.head(100)
    total_rows = len(df)
    results = []

    for idx, row in tqdm(df.iterrows(), total=len(df), desc="Processing rows"):
        premise = row['post_body']  # Assuming 'post_body' column exists
        probabilities, sentence_level_results = classify_text(premise, sent=True)
        
        # If no results, continue to the next row
        if probabilities is None or len(probabilities) == 0 or sentence_level_results is None:
            continue
        
        # Extract sentiment separately so it does not compete with classifications
        sentiment_score = probabilities.get('sentiment_score')
        classification_probs = {
            label: prob for label, prob in probabilities.items()
            if label != 'sentiment_score'
        }
        
        # Parse probabilities into following columns
        sorted_probs = sorted(classification_probs.items(), key=lambda item: item[1], reverse=True)
        
        primary_label = sorted_probs[0][0]  # Label with the highest probability
        primary_probability = sorted_probs[0][1]
        secondary_labels = [label for label, _ in sorted_probs[1:]]  # Remainder labels
        secondary_probabilities = [prob for _, prob in sorted_probs[1:]]
        # all_post_level_probabilities = [prob for _, prob in sorted_probabilities]  # All probabilities

        # Prepare row result
        row_result = {
            'Post': premise,
            'Post sentiment score': round(float(sentiment_score), 3),
            'Primary livability classification': primary_label,
            'Primary classification confidence': format_probability(primary_probability),
            'Secondary livability classifications': ', '.join(secondary_labels),
            'Secondary classification confidences': ', '.join(format_probability(prob) for prob in secondary_probabilities),
            'Sentence-level results': format_sentence_level_results(sentence_level_results)
        }
        results.append(row_result)
        
        # Simulate processing time and update progress
        progress = f"Processing rows: {idx}/{total_rows}"
        yield progress, None, None

    # Convert results to DataFrame
    result_df = pd.DataFrame(results)
    result_df_sliced = result_df.head(5)
    result_df_sliced = result_df_sliced.drop(columns=['Sentence-level results'])
    
    nan_or_empty = result_df_sliced['Secondary livability classifications'].map(
        lambda x: x == '' or pd.isna(x)
    ).all()
    
    if nan_or_empty:
        result_df_sliced = result_df_sliced.drop(columns=['Secondary livability classifications', 'Secondary classification confidences'])

    result_df_sliced = result_df_sliced[
        [column for column in result_df_sliced.columns if column != 'Post sentiment score']
        + ['Post sentiment score']
    ]

    # Write the CSV to a temporary file
    with NamedTemporaryFile(delete=False, suffix=".csv") as temp_file:
        result_df.to_csv(temp_file.name, index=False)
        temp_file_path = temp_file.name

    if user_consent:
        try:
            # Get current time in EST
            est = pytz.timezone("US/Eastern")
            timestamp = datetime.now(est).strftime("%Y-%m-%d_%H-%M-%S")
            if email_input is not None:
                filename = f"classified_data_{email_input}_{timestamp}.csv"
            else:
                filename = f"classified_data_{timestamp}.csv"
                
            upload_result = upload_to_hf_dataset(temp_file_path, f"processed/{filename}")
            print(f"Classified data saved to Hugging Face: {upload_result}")
        except Exception as e:
            print("Error saving classified data to hf:", e)

    yield "Processing complete!", result_df_sliced, temp_file_path
    return


# Gradio app
with gr.Blocks(title="Housing Livability Classifier") as demo:
    gr.Markdown("""
    ---
    [Checkout](https://github.com/Herb-Lab/LLM_housing_livability) our GitHub repository
    [View](https://huggingface.co/Herb-Lab/LLM_housing_livability) the model
    """)
    with gr.Tabs():
        with gr.TabItem("Text Classification"):
            gr.Markdown("## Housing Livability Text Classifier")
            gr.Markdown("Classify text into Indoor Air Quality, Thermal, Acoustic, or Visual categories.")

            text_input = gr.Textbox(lines=5, placeholder="Enter text to classify...")
            post_level_output = gr.Label(num_top_classes=4, label="Post Level Results")
            sentence_level_output = gr.Dataframe(label="Sentence Level Results",
                                                 column_widths=["20%", "10%", "15%", "15%", "15%"])

            submit_button = gr.Button("Classify")
            submit_button.click(fn=classify_text_for_display,
                                inputs=text_input,
                                outputs=[post_level_output, sentence_level_output])

        with gr.TabItem("CSV Bulk Classification"):
            gr.Markdown("## CSV Bulk Classification")
            gr.Markdown("Upload a CSV file containing a single column. The first row should be a header (it will be ignored during processing)."
                        "Only the first 100 rows will be processed. If you wish to process more than 100 posts, please contact the authors."
                        "The app will classify each entry and generate a downloadable CSV with the results.")
            
            gr.Markdown("### ⚠️ Consent Required")
            gr.Markdown("**Please confirm your consent to proceed. Your input reviews and the output may be shared with the authors.**")
            consent_checkbox = gr.Checkbox(
                label="✅ I consent to sharing my input posts and output with the authors (Mandatory)", value=False
            )
            
            with gr.Row():
                file_input = gr.File(label="Upload CSV File")
                
            with gr.Column():
                # email_input = gr.Textbox(label="Email (optional, if you want to receive processed results by email)")
                submit_button = gr.Button("Submit", interactive=False)  # Initially disabled

            clear_button = gr.Button("Clear")

            progress_output = gr.Textbox(label="Progress")

            top_results_output = gr.Dataframe(label="Top 5 Results",
                                            column_widths=["20%", "10%", "10%", "10%", "10%", "20%"],
                                            scale=100)

            download_button = gr.File(label="Download Processed CSV")

            # Enable/disable submit button based on checkbox
            def toggle_submit(consent):
                return gr.update(interactive=consent)

            consent_checkbox.change(toggle_submit, inputs=consent_checkbox, outputs=submit_button)

            submit_button.click(
                fn=process_csv,
                inputs=[file_input, consent_checkbox],
                outputs=[progress_output, top_results_output, download_button]
            )
demo.queue()
demo.launch()