Update app.py

app.py

@@ -1,21 +1,437 @@
-import warnings; warnings.filterwarnings("ignore", message="pkg_resources is deprecated")
 import gradio as gr
-from pipeline import run_pipeline
-import gradio_client, fastapi
-print("VERSIONS -> gradio:", gr.__version__, "gradio_client:", gradio_client.__version__, "fastapi:", fastapi.__version__)
-
-def infer(file):
-    path = getattr(file, "name", file)  # gr.File gives a tempfile
-    return run_pipeline(path, out_style="json")  # returns JSON string
-
-demo = gr.Interface(
-    fn=infer,
-    inputs=gr.File(label="Upload audio/video (mp3, mp4, wav)"),
-    outputs=gr.Textbox(label="Result (JSON text)"),  # ← was gr.JSON; use Textbox
-    title="Aphasia Classification",
-    description="MP3/MP4 → WAV → .cha → JSON → model",
-    concurrency_limit=1,
-)
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import json
+import numpy as np
+import math
+from transformers import AutoTokenizer, AutoModel
+from typing import Dict, List, Optional
+import logging
 
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Recreate the model classes (simplified versions)
+class StablePositionalEncoding(nn.Module):
+    def __init__(self, d_model: int, max_len: int = 5000):
+        super().__init__()
+        self.d_model = d_model
+        
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * 
+                           (-math.log(10000.0) / d_model))
+        
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        
+        self.register_buffer('pe', pe.unsqueeze(0))
+        self.learnable_pe = nn.Parameter(torch.randn(max_len, d_model) * 0.01)
+    
+    def forward(self, x):
+        seq_len = x.size(1)
+        sinusoidal = self.pe[:, :seq_len, :].to(x.device)
+        learnable = self.learnable_pe[:seq_len, :].unsqueeze(0).expand(x.size(0), -1, -1)
+        return x + 0.1 * (sinusoidal + learnable)
+
+class StableMultiHeadAttention(nn.Module):
+    def __init__(self, feature_dim: int, num_heads: int = 4, dropout: float = 0.3):
+        super().__init__()
+        self.num_heads = num_heads
+        self.feature_dim = feature_dim
+        self.head_dim = feature_dim // num_heads
+        
+        assert feature_dim % num_heads == 0
+        
+        self.query = nn.Linear(feature_dim, feature_dim)
+        self.key = nn.Linear(feature_dim, feature_dim)
+        self.value = nn.Linear(feature_dim, feature_dim)
+        self.dropout = nn.Dropout(dropout)
+        self.output_proj = nn.Linear(feature_dim, feature_dim)
+        self.layer_norm = nn.LayerNorm(feature_dim)
+        
+    def forward(self, x, mask=None):
+        batch_size, seq_len, _ = x.size()
+        
+        Q = self.query(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+        K = self.key(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+        V = self.value(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+        
+        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.head_dim)
+        
+        if mask is not None:
+            if mask.dim() == 2:
+                mask = mask.unsqueeze(1).unsqueeze(1)
+            scores.masked_fill_(mask == 0, -1e9)
+        
+        attn_weights = F.softmax(scores, dim=-1)
+        attn_weights = self.dropout(attn_weights)
+        
+        context = torch.matmul(attn_weights, V)
+        context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.feature_dim)
+        
+        output = self.output_proj(context)
+        return self.layer_norm(output + x)
+
+class StableLinguisticFeatureExtractor(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        
+        # Simplified version - just return zeros if features not available
+        self.pos_embedding = nn.Embedding(config.get('pos_vocab_size', 150), config.get('pos_emb_dim', 64), padding_idx=0)
+        self.pos_attention = StableMultiHeadAttention(config.get('pos_emb_dim', 64), num_heads=4)
+        
+        self.grammar_projection = nn.Sequential(
+            nn.Linear(config.get('grammar_dim', 3), config.get('grammar_hidden_dim', 64)),
+            nn.Tanh(),
+            nn.LayerNorm(config.get('grammar_hidden_dim', 64)),
+            nn.Dropout(config.get('dropout_rate', 0.3) * 0.3)
+        )
+        
+        self.duration_projection = nn.Sequential(
+            nn.Linear(1, config.get('duration_hidden_dim', 128)),
+            nn.Tanh(),
+            nn.LayerNorm(config.get('duration_hidden_dim', 128))
+        )
+        
+        self.prosody_projection = nn.Sequential(
+            nn.Linear(config.get('prosody_dim', 32), config.get('prosody_dim', 32)),
+            nn.ReLU(),
+            nn.LayerNorm(config.get('prosody_dim', 32))
+        )
+        
+        total_feature_dim = (config.get('pos_emb_dim', 64) + config.get('grammar_hidden_dim', 64) + 
+                           config.get('duration_hidden_dim', 128) + config.get('prosody_dim', 32))
+        self.feature_fusion = nn.Sequential(
+            nn.Linear(total_feature_dim, total_feature_dim // 2),
+            nn.Tanh(),
+            nn.LayerNorm(total_feature_dim // 2),
+            nn.Dropout(config.get('dropout_rate', 0.3))
+        )
+        
+    def forward(self, pos_ids, grammar_ids, durations, prosody_features, attention_mask):
+        batch_size, seq_len = pos_ids.size()
+        
+        # Simple processing - can be expanded later
+        pos_ids_clamped = pos_ids.clamp(0, self.config.get('pos_vocab_size', 150) - 1)
+        pos_embeds = self.pos_embedding(pos_ids_clamped)
+        pos_features = self.pos_attention(pos_embeds, attention_mask)
+        
+        grammar_features = self.grammar_projection(grammar_ids.float())
+        duration_features = self.duration_projection(durations.unsqueeze(-1).float())
+        prosody_features = self.prosody_projection(prosody_features.float())
+        
+        combined_features = torch.cat([
+            pos_features, grammar_features, duration_features, prosody_features
+        ], dim=-1)
+        
+        fused_features = self.feature_fusion(combined_features)
+        
+        mask_expanded = attention_mask.unsqueeze(-1).float()
+        pooled_features = torch.sum(fused_features * mask_expanded, dim=1) / torch.sum(mask_expanded, dim=1)
+        
+        return pooled_features
+
+class StableAphasiaClassifier(nn.Module):
+    def __init__(self, config, num_labels: int):
+        super().__init__()
+        self.config = config
+        self.num_labels = num_labels
+        
+        try:
+            # Load the base BERT model
+            self.bert = AutoModel.from_pretrained(config.get('model_name', 'microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext'))
+            self.bert_config = self.bert.config
+            
+            # Freeze embeddings for stability
+            for param in self.bert.embeddings.parameters():
+                param.requires_grad = False
+            
+            self.positional_encoder = StablePositionalEncoding(
+                d_model=self.bert_config.hidden_size,
+                max_len=config.get('max_length', 512)
+            )
+            
+            self.linguistic_extractor = StableLinguisticFeatureExtractor(config)
+            
+            bert_dim = self.bert_config.hidden_size
+            linguistic_dim = (config.get('pos_emb_dim', 64) + config.get('grammar_hidden_dim', 64) + 
+                             config.get('duration_hidden_dim', 128) + config.get('prosody_dim', 32)) // 2
+            
+            self.feature_fusion = nn.Sequential(
+                nn.Linear(bert_dim + linguistic_dim, bert_dim),
+                nn.LayerNorm(bert_dim),
+                nn.Tanh(),
+                nn.Dropout(config.get('dropout_rate', 0.3))
+            )
+            
+            # Classifier
+            self.classifier = self._build_classifier(bert_dim, num_labels, config)
+            
+            # Multi-task heads (simplified)
+            self.severity_head = nn.Sequential(
+                nn.Linear(bert_dim, 4),
+                nn.Softmax(dim=-1)
+            )
+            
+            self.fluency_head = nn.Sequential(
+                nn.Linear(bert_dim, 1),
+                nn.Sigmoid()
+            )
+            
+        except Exception as e:
+            logger.error(f"Error initializing model: {e}")
+            raise
+        
+    def _build_classifier(self, input_dim: int, num_labels: int, config):
+        layers = []
+        current_dim = input_dim
+        
+        classifier_hidden_dims = config.get('classifier_hidden_dims', [512, 256])
+        for hidden_dim in classifier_hidden_dims:
+            layers.extend([
+                nn.Linear(current_dim, hidden_dim),
+                nn.LayerNorm(hidden_dim),
+                nn.Tanh(),
+                nn.Dropout(config.get('dropout_rate', 0.3))
+            ])
+            current_dim = hidden_dim
+        
+        layers.append(nn.Linear(current_dim, num_labels))
+        return nn.Sequential(*layers)
+    
+    def forward(self, input_ids, attention_mask, labels=None,
+                word_pos_ids=None, word_grammar_ids=None, word_durations=None,
+                prosody_features=None, **kwargs):
+        
+        bert_outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
+        sequence_output = bert_outputs.last_hidden_state
+        
+        position_enhanced = self.positional_encoder(sequence_output)
+        pooled_output = self._attention_pooling(position_enhanced, attention_mask)
+        
+        # Handle missing linguistic features
+        if all(x is not None for x in [word_pos_ids, word_grammar_ids, word_durations]):
+            if prosody_features is None:
+                batch_size, seq_len = input_ids.size()
+                prosody_features = torch.zeros(
+                    batch_size, seq_len, self.config.get('prosody_dim', 32),
+                    device=input_ids.device
+                )
+            
+            linguistic_features = self.linguistic_extractor(
+                word_pos_ids, word_grammar_ids, word_durations,
+                prosody_features, attention_mask
+            )
+        else:
+            # Create dummy linguistic features
+            linguistic_features = torch.zeros(
+                input_ids.size(0), 
+                (self.config.get('pos_emb_dim', 64) + self.config.get('grammar_hidden_dim', 64) + 
+                 self.config.get('duration_hidden_dim', 128) + self.config.get('prosody_dim', 32)) // 2,
+                device=input_ids.device
+            )
+        
+        combined_features = torch.cat([pooled_output, linguistic_features], dim=1)
+        fused_features = self.feature_fusion(combined_features)
+        
+        logits = self.classifier(fused_features)
+        severity_pred = self.severity_head(fused_features)
+        fluency_pred = self.fluency_head(fused_features)
+        
+        return {
+            "logits": logits,
+            "severity_pred": severity_pred,
+            "fluency_pred": fluency_pred,
+        }
+    
+    def _attention_pooling(self, sequence_output, attention_mask):
+        attention_weights = torch.softmax(
+            torch.sum(sequence_output, dim=-1, keepdim=True), dim=1
+        )
+        attention_weights = attention_weights * attention_mask.unsqueeze(-1).float()
+        attention_weights = attention_weights / (torch.sum(attention_weights, dim=1, keepdim=True) + 1e-9)
+        pooled = torch.sum(sequence_output * attention_weights, dim=1)
+        return pooled
+
+# Load configuration and model
+def load_model():
+    try:
+        # Load configuration
+        with open("config.json", "r") as f:
+            config = json.load(f)
+        
+        logger.info(f"Loaded config: {config}")
+        
+        # Initialize tokenizer
+        model_name = config.get('model_name', 'microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext')
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        
+        # Add special tokens
+        special_tokens = ["[DIALOGUE]", "[TURN]", "[PAUSE]", "[REPEAT]", "[HESITATION]"]
+        tokenizer.add_special_tokens({"additional_special_tokens": special_tokens})
+        
+        # Initialize model
+        num_labels = config.get('num_labels', 9)
+        model_config = config.get('model_config', {})
+        
+        model = StableAphasiaClassifier(model_config, num_labels)
+        model.bert.resize_token_embeddings(len(tokenizer))
+        
+        # Load model weights
+        try:
+            state_dict = torch.load("pytorch_model.bin", map_location="cpu")
+            model.load_state_dict(state_dict)
+            logger.info("Successfully loaded model weights")
+        except Exception as e:
+            logger.error(f"Error loading model weights: {e}")
+            logger.info("Using randomly initialized weights")
+        
+        model.eval()
+        
+        # Get label mapping
+        id2label = config.get('id2label', {})
+        
+        return model, tokenizer, id2label
+        
+    except Exception as e:
+        logger.error(f"Error loading model: {e}")
+        raise
+
+# Initialize model (with error handling)
+try:
+    model, tokenizer, id2label = load_model()
+    logger.info("Model loaded successfully!")
+except Exception as e:
+    logger.error(f"Failed to load model: {e}")
+    model, tokenizer, id2label = None, None, {}
+
+def predict_aphasia(text):
+    """Predict aphasia type from text"""
+    try:
+        if model is None or tokenizer is None:
+            return "Error: Model not loaded properly. Please check the logs.", 0.0, "N/A", 0.0
+        
+        if not text or not text.strip():
+            return "Please enter some text for analysis.", 0.0, "N/A", 0.0
+        
+        # Tokenize input
+        inputs = tokenizer(
+            text,
+            max_length=512,
+            padding="max_length",
+            truncation=True,
+            return_tensors="pt"
+        )
+        
+        # Create dummy linguistic features (since we don't have them from raw text)
+        batch_size, seq_len = inputs["input_ids"].size()
+        dummy_pos = torch.zeros(batch_size, seq_len, dtype=torch.long)
+        dummy_grammar = torch.zeros(batch_size, seq_len, 3, dtype=torch.long)
+        dummy_durations = torch.zeros(batch_size, seq_len, dtype=torch.float)
+        dummy_prosody = torch.zeros(batch_size, seq_len, 32, dtype=torch.float)
+        
+        # Make prediction
+        with torch.no_grad():
+            outputs = model(
+                input_ids=inputs["input_ids"],
+                attention_mask=inputs["attention_mask"],
+                word_pos_ids=dummy_pos,
+                word_grammar_ids=dummy_grammar,
+                word_durations=dummy_durations,
+                prosody_features=dummy_prosody
+            )
+        
+        # Process outputs
+        logits = outputs["logits"]
+        probs = F.softmax(logits, dim=1)
+        predicted_class_id = torch.argmax(probs, dim=1).item()
+        confidence = torch.max(probs, dim=1)[0].item()
+        
+        # Get predicted label
+        predicted_label = id2label.get(str(predicted_class_id), f"Class_{predicted_class_id}")
+        
+        # Get additional predictions
+        severity = torch.argmax(outputs["severity_pred"], dim=1).item()
+        fluency = outputs["fluency_pred"].item()
+        
+        # Format results
+        result = f"Predicted Aphasia Type: {predicted_label}"
+        confidence_str = f"Confidence: {confidence:.2%}"
+        severity_str = f"Severity Level: {severity}/3"
+        fluency_str = f"Fluency Score: {fluency:.3f}"
+        
+        return result, confidence, severity_str, fluency_str
+        
+    except Exception as e:
+        logger.error(f"Prediction error: {e}")
+        return f"Error during prediction: {str(e)}", 0.0, "N/A", 0.0
+
+# Create Gradio interface
+def create_interface():
+    """Create Gradio interface"""
+    
+    with gr.Blocks(title="Aphasia Classification System") as demo:
+        gr.Markdown("# 🧠 Advanced Aphasia Classification System")
+        gr.Markdown("Enter speech or text data to classify aphasia type and analyze linguistic patterns.")
+        
+        with gr.Row():
+            with gr.Column():
+                text_input = gr.Textbox(
+                    label="Input Text",
+                    placeholder="Enter speech transcription or text for analysis...",
+                    lines=5
+                )
+                
+                submit_btn = gr.Button("Analyze Text", variant="primary")
+                clear_btn = gr.Button("Clear", variant="secondary")
+                
+            with gr.Column():
+                prediction_output = gr.Textbox(label="Prediction Result", lines=2)
+                confidence_output = gr.Textbox(label="Confidence Score", lines=1)
+                severity_output = gr.Textbox(label="Severity Analysis", lines=1)
+                fluency_output = gr.Textbox(label="Fluency Analysis", lines=1)
+        
+        # Event handlers
+        submit_btn.click(
+            fn=predict_aphasia,
+            inputs=[text_input],
+            outputs=[prediction_output, confidence_output, severity_output, fluency_output]
+        )
+        
+        clear_btn.click(
+            fn=lambda: ("", "", "", "", ""),
+            inputs=[],
+            outputs=[text_input, prediction_output, confidence_output, severity_output, fluency_output]
+        )
+        
+        # Add examples
+        gr.Examples(
+            examples=[
+                ["The patient... uh... wants to... go home but... cannot... find the words"],
+                ["Woman is... is washing dishes and the... the... sink is overflowing with water everywhere"],
+                ["Cookie is in the cookie jar on the... on the... what do you call it... the shelf thing"]
+            ],
+            inputs=text_input
+        )
+        
+        gr.Markdown("### About")
+        gr.Markdown("This system uses a specialized transformer model trained on clinical speech data to classify different types of aphasia.")
+    
+    return demo
+
+# Launch the app
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True, max_threads=1, share=True)  # ← share=True
+    try:
+        demo = create_interface()
+        demo.launch(
+            server_name="0.0.0.0",
+            server_port=7860,
+            show_error=True
+        )
+    except Exception as e:
+        logger.error(f"Failed to launch app: {e}")
+        print(f"Application startup failed: {e}")