Update app.py

app.py

@@ -1,181 +1,182 @@
-# 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)
-
-    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, fed into a Multi-Layer Perceptron (MLP) "
-    "to classify text as either 'Normal Reading (NR)' or 'Task-Specific Reading (TSR)' based on the ZuCo dataset."
-)
-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').
+# 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, fed into a Multi-Layer Perceptron (MLP) "
+    "to classify text as either 'Normal Reading (NR)' or 'Task-Specific Reading (TSR)' based on the ZuCo dataset."
+)
+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)