Update app.py

app.py

@@ -1,33 +1,58 @@
 import gradio as gr
 import mne
+import numpy as np
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
-# Load open-source LLM (no training needed)
+# Load an open-source LLM model with no additional training
 model_name = "tiiuae/falcon-7b-instruct"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, torch_dtype=torch.float16, device_map="auto")
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    trust_remote_code=True,
+    torch_dtype=torch.float16,
+    device_map="auto"  # Automatically selects CPU/GPU if available
+)
+
+def compute_band_power(psd, freqs, fmin, fmax):
+    """Compute mean band power in the given frequency range."""
+    freq_mask = (freqs >= fmin) & (freqs <= fmax)
+    # Take the mean across channels and frequencies
+    band_psd = psd[:, freq_mask].mean()  
+    return float(band_psd)
 
 def process_eeg(file):
     # Load EEG data using MNE
+    # This expects a .fif file containing raw EEG data
     raw = mne.io.read_raw_fif(file.name, preload=True)
-    # Compute some features (e.g., average band powers)
+
+    # Compute PSD (Power Spectral Density) between 1 and 40 Hz
     psd, freqs = mne.time_frequency.psd_welch(raw, fmin=1, fmax=40)
+
+    # Compute simple band powers
     alpha_power = compute_band_power(psd, freqs, 8, 12)
     beta_power = compute_band_power(psd, freqs, 13, 30)
-    # Create a human-readable summary of features
-    data_summary = f"Alpha power: {alpha_power}, Beta power: {beta_power}. The data shows stable alpha rhythms and slightly elevated beta."
-    
-    # Prompt the LLM
-    prompt = f"""You are a neuroscientist analyzing EEG features. 
-    Data Summary: {data_summary}
-
-    Provide a concise, user-friendly interpretation of these findings."""
-    
+
+    # Create a short summary of the extracted features
+    data_summary = (
+        f"Alpha power: {alpha_power:.3f}, Beta power: {beta_power:.3f}. "
+        f"The EEG shows stable alpha rhythms and slightly elevated beta activity."
+    )
+
+    # Prepare the prompt for the language model
+    prompt = f"""You are a neuroscientist analyzing EEG features.
+Data Summary: {data_summary}
+
+Provide a concise, user-friendly interpretation of these findings in simple terms.
+"""
+
+    # Generate the summary using the LLM
     inputs = tokenizer.encode(prompt, return_tensors="pt").to(model.device)
-    outputs = model.generate(inputs, max_length=200, do_sample=True, top_k=50, top_p=0.95)
+    outputs = model.generate(
+        inputs, max_length=200, do_sample=True, top_k=50, top_p=0.95
+    )
     summary = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    
+
     return summary
 
 iface = gr.Interface(
@@ -35,7 +60,7 @@ iface = gr.Interface(
     inputs=gr.File(label="Upload your EEG data (FIF format)"),
     outputs="text",
     title="NeuroNarrative-Lite: EEG Summary",
-    description="Upload EEG data to receive a text-based summary from an open-source LLM. No training required!"
+    description="Upload EEG data to receive a text-based summary from an open-source language model. No training required!"
 )
 
 if __name__ == "__main__":