app.py

# app.py
import gradio as gr
import torch
import torch.nn as nn
import numpy as np
import pandas as pd
import spacy
import textstat
from nltk.tokenize import word_tokenize
import nltk
import re
import joblib
from transformers import BertTokenizerFast, BertForSequenceClassification
from sentence_transformers import SentenceTransformer

# --- 1. SETUP: Constants and Model Loading ---
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
FINETUNE_MODEL_NAME = 'bert-base-uncased'
MAX_LEN_BERT = 128
print(f"Using device: {DEVICE}")
NLP = spacy.load('en_core_web_sm', disable=['ner'])
SCALER = joblib.load('scaler_mlp_discrete.joblib')

# --- (Re)Define the PyTorch MLP Model Class ---
class AdvancedMLP(nn.Module):
    # ... (This class is correct, no changes needed)
    def __init__(self, input_dim, num_classes=2):
        super(AdvancedMLP, self).__init__()
        self.layer_1 = nn.Linear(input_dim, 512)
        self.relu1 = nn.ReLU()
        self.batchnorm1 = nn.BatchNorm1d(512)
        self.dropout1 = nn.Dropout(0.3)
        self.layer_2 = nn.Linear(512, 128)
        self.relu2 = nn.ReLU()
        self.batchnorm2 = nn.BatchNorm1d(128)
        self.dropout2 = nn.Dropout(0.3)
        self.output_layer = nn.Linear(128, num_classes)

    def forward(self, x):
        x = self.layer_1(x); x = self.relu1(x); x = self.batchnorm1(x); x = self.dropout1(x)
        x = self.layer_2(x); x = self.relu2(x); x = self.batchnorm2(x); x = self.dropout2(x)
        x = self.output_layer(x)
        return x

# --- Load All Models and Artifacts ---
print("Loading models and artifacts...")
try:
    nltk.download('punkt', quiet=True)
    nltk.download('punkt_tab', quiet=True)

    TOKENIZER = BertTokenizerFast.from_pretrained(FINETUNE_MODEL_NAME)
    
    bert_for_seq_clf = BertForSequenceClassification.from_pretrained(FINETUNE_MODEL_NAME, num_labels=2)
    # NOTE: Ensure you have the correct file for the best BERT model. The user provided 'fold_4'.
    bert_for_seq_clf.load_state_dict(torch.load("best_bert_finetuned_fold_4.bin", map_location=DEVICE))
    BERT_EMBEDDING_MODEL = bert_for_seq_clf.bert.to(DEVICE).eval()

    INPUT_DIM_MLP = 768 + 19 
    MLP_MODEL = AdvancedMLP(input_dim=INPUT_DIM_MLP).to(DEVICE)
    MLP_MODEL.load_state_dict(torch.load("best_mlp_combined_features_ZuCo.bin", map_location=DEVICE))
    MLP_MODEL.eval()

    NLP = spacy.load('en_core_web_sm', disable=['ner'])

    # NOTE: Ensure this filename matches the scaler you saved.
    SCALER = joblib.load('scaler_mlp_discrete.joblib')
    
    print("All models and artifacts loaded successfully.")

except FileNotFoundError as e:
    print(f"ERROR: A required file was not found: {e.name}")
    print("Please ensure 'best_bert_finetuned_fold_4.bin', 'best_mlp_combined_features_ZuCo.bin', and 'scaler_mlp_discrete.joblib' are in the same directory.")
    exit()

# --- 2. PREPROCESSING & FEATURE ENGINEERING FUNCTIONS ---
def clean_text(text):
    text = str(text).lower()
    return re.sub(r'\\s+', ' ', text).strip()

# FIX 1: Pass the `nlp_model` object as an argument.
def get_discrete_features(sentence, nlp_model):
    """Calculates all 19 discrete linguistic features for a single sentence."""
    features = {}
    
    # ... (rest of the feature calculation is correct)
    features['char_count'] = len(sentence)
    words = sentence.split()
    features['word_count'] = len(words)
    features['avg_word_length'] = features['char_count'] / features['word_count'] if features['word_count'] > 0 else 0
    features['flesch_ease'] = textstat.flesch_reading_ease(sentence)
    features['flesch_grade'] = textstat.flesch_kincaid_grade(sentence)
    features['gunning_fog'] = textstat.gunning_fog(sentence)
    tokens = word_tokenize(sentence)
    features['ttr'] = len(set(tokens)) / len(tokens) if tokens else 0
    features['lex_density_proxy'] = sum(1 for w in tokens if len(w) > 6) / len(tokens) if tokens else 0
    
    # FIX 2: Use the passed `nlp_model` argument instead of the global name `NLP`.
    doc = nlp_model(sentence)
    dep_distances = [abs(token.i - token.head.i) for token in doc if token.head is not token]
    pos_counts = doc.count_by(spacy.attrs.POS)
    
    features['num_subord_clauses'] = sum(1 for token in doc if token.dep_ == 'mark')
    features['num_conj_clauses'] = sum(1 for token in doc if token.dep_ == 'cc' and token.head.pos_ == 'VERB')
    features['avg_dep_dist'] = np.mean(dep_distances) if dep_distances else 0
    features['max_dep_dist'] = np.max(dep_distances) if dep_distances else 0
    features['num_verbs'] = pos_counts.get(spacy.parts_of_speech.VERB, 0)
    features['num_nouns'] = pos_counts.get(spacy.parts_of_speech.NOUN, 0) + pos_counts.get(spacy.parts_of_speech.PROPN, 0)
    features['num_adjectives'] = pos_counts.get(spacy.parts_of_speech.ADJ, 0)
    features['num_adverbs'] = pos_counts.get(spacy.parts_of_speech.ADV, 0)
    features['num_prepositions'] = pos_counts.get(spacy.parts_of_speech.ADP, 0)
    features['num_conjunctions'] = pos_counts.get(spacy.parts_of_speech.CCONJ, 0) + pos_counts.get(spacy.parts_of_speech.SCONJ, 0)
    
    feature_order = [
        'char_count', 'word_count', 'avg_word_length', 'ttr', 'lex_density_proxy',
        'flesch_ease', 'flesch_grade', 'gunning_fog', 'num_subord_clauses',
        'num_conj_clauses', 'avg_dep_dist', 'max_dep_dist', 'num_verbs',
        'num_nouns', 'num_adjectives', 'num_adverbs', 'num_prepositions', 'num_conjunctions',
        'ollama_llm_rating'
    ]
    features['ollama_llm_rating'] = 3.0
    return np.array([features[k] for k in feature_order]).reshape(1, -1)

def get_bert_embedding(sentence):
    # ... (This function is correct, no changes needed)
    encoded = TOKENIZER.encode_plus(sentence, add_special_tokens=True, max_length=MAX_LEN_BERT, return_token_type_ids=False, padding='max_length', truncation=True, return_attention_mask=True, return_tensors='pt')
    input_ids, attention_mask = encoded['input_ids'].to(DEVICE), encoded['attention_mask'].to(DEVICE)
    with torch.no_grad():
        outputs = BERT_EMBEDDING_MODEL(input_ids, attention_mask=attention_mask)
        embedding = outputs.last_hidden_state[:, 0, :].cpu().numpy()
    return embedding

# --- 3. THE PREDICTION FUNCTION ---
def predict_cognitive_state(sentence):
    if not sentence.strip():
        return {"Normal Reading (NR)": 0, "Task-Specific Reading (TSR)": 0}

    cleaned = clean_text(sentence)

    # FIX 3: Pass the loaded NLP model into the function.
    discrete_features = get_discrete_features(cleaned, NLP)
    
    scaled_discrete_features = SCALER.transform(discrete_features)
    bert_embedding = get_bert_embedding(cleaned)
    combined_features = np.concatenate((bert_embedding, scaled_discrete_features), axis=1)
    
    features_tensor = torch.tensor(combined_features, dtype=torch.float32).to(DEVICE)
    with torch.no_grad():
        logits = MLP_MODEL(features_tensor)
        probabilities = torch.softmax(logits, dim=1).cpu().numpy()[0]

    labels = ["Normal Reading (NR)", "Task-Specific Reading (TSR)"]
    confidences = {label: float(prob) for label, prob in zip(labels, probabilities)}
    
    return confidences

# --- 4. GRADIO INTERFACE ---
title = "🧠 Cognitive State Analysis from Text"
description = (
    "Enter a sentence to predict its cognitive state. This demo uses a fine-tuned BERT model for semantic "
    "embeddings combined with 19 discrete linguistic features. These features are fed into a Multi-Layer Perceptron (MLP) "
    "to classify the input as either:\n\n"
    "- **Normal Reading (NR):** Casual reading without a specific goal—like reading a story or browsing news.\n"
    "- **Task-Specific Reading (TSR):** Purpose-driven reading—such as searching for an answer or following instructions.\n\n"
    "The model is trained on text data from the ZuCo dataset, using only linguistic features—no EEG or eye-tracking signals are used."
)
example_list = [
    ["Through his son Timothy Bush, Jr., who was also a blacksmith, descended two American Presidents -George H. W. Bush and George W. Bush."],
    ["He received his bachelor's degree in 1965 and master's degree in political science in 1966 both from the University of Wyoming."],
    ["What does the abbreviation Ph.D. stand for?"],
    ["What is the name of the director of the 2003 American film 'The Haunted Mansion'?"],
]

demo = gr.Interface(
    fn=predict_cognitive_state,
    inputs=gr.Textbox(lines=3, label="Input Sentence", placeholder="Type a sentence here..."),
    outputs=gr.Label(num_top_classes=2, label="Prediction"),
    title=title,
    description=description,
    examples=example_list,
    allow_flagging="never"
)

if __name__ == "__main__":
    # FIX 4: Corrected the typo from Launch to launch (lowercase 'l').
    demo.launch(debug=True)