Added gradio demo interface

scripts/demo_gradio.py

@@ -0,0 +1,425 @@
+"""
+Gradio Demo interface for LexiMind NLP pipeline.
+Showcases summarization, emotion detection, and topic prediction.
+"""
+import json
+import sys
+from io import StringIO
+from pathlib import Path
+from typing import Iterable, Sequence
+import gradio as gr
+from gradio.themes import Soft
+import matplotlib.pyplot as plt
+import pandas as pd
+import seaborn as sns
+import torch
+from matplotlib.figure import Figure
+
+# Add project root to the path, going up two folder levels from this file
+project_root = Path(__file__).parent.parent
+sys.path.insert(0, str(project_root))
+
+from src.inference.factory import create_inference_pipeline
+from src.inference.pipeline import EmotionPrediction, InferencePipeline, TopicPrediction
+from src.utils.logging import configure_logging, get_logger
+
+configure_logging()
+logger = get_logger(__name__)
+
+_pipeline: InferencePipeline | None = None  # Global pipeline instance
+_label_metadata = None  # Cached label metadata
+
+
+def get_pipeline() -> InferencePipeline:
+    """Lazy Loading and Caching the inference pipeline"""
+    global _pipeline, _label_metadata
+    if _pipeline is None:
+        try:
+            logger.info("Loading inference pipeline...")
+            pipeline, label_metadata = create_inference_pipeline(
+                tokenizer_dir="data/tokenization",
+                checkpoint_path="checkpoints/best.pt",
+                labels_path="data/labels.json",
+            )
+            _pipeline = pipeline
+            _label_metadata = label_metadata
+            logger.info("Pipeline loaded successfully")
+        except Exception as e:
+            logger.error(f"Failed to load pipeline: {e}")
+            raise RuntimeError("Could not initialize inference pipeline. Check logs for details.")
+    return _pipeline
+
+def count_tokens(text: str) -> str:
+    """Count tokens in the input text."""
+    if not text:
+        return "Tokens: 0"
+    try: 
+        pipeline = get_pipeline()
+        token_count = len(pipeline.tokenizer.encode(text))
+        return f"Tokens: {token_count}"
+    except Exception as e:
+        logger.error(f"Token counting error: {e}")
+        return "Token count unavailable"
+            
+def map_compression_to_length(compression: int, max_model_length: int = 512):
+    """
+    Map Compression slider (20-80%) to max summary length.
+    Higher compression = shorter summary output.
+    """
+    # Invert, 20% compression = 80% of max length
+    ratio = (100 - compression) / 100
+    return int(ratio * max_model_length)
+
+def predict(text: str, compression: int):
+    """
+    Main predcition function for the Gradio interface.
+    Args:
+        text: Text to process
+        compression: Compression percentage (20-80)
+    Returns:
+        Tuple of (summary_html, emotion_plot, topic_output, attention_fig, download_data)
+    """
+    if not text or not text.strip():
+        return ("Please enter some text to analyze.", 
+                None, 
+                "No topic prediction available", 
+                None, 
+                None)
+    try:
+        pipeline = get_pipeline()
+        max_len = map_compression_to_length(compression)
+        logger.info(f"Generating summary with max length of {max_len}")
+        
+        # Get the predictions
+        summary = pipeline.summarize([text], max_length=max_len)[0]
+        emotions = pipeline.predict_emotions([text])[0]
+        topic = pipeline.predict_topics([text])[0]
+        
+        summary_html = format_summary(text, summary)
+        emotion_plot = create_emotion_plot(emotions)
+        topic_output = format_topic(topic)
+        attention_fig = create_attention_heatmap(text, summary, pipeline)
+        download_data = prepare_download(text, summary, emotions, topic)
+        
+        return summary_html, emotion_plot, topic_output, attention_fig, download_data
+    
+    except Exception as e:
+        logger.error(f"Prediction error: {e}", exc_info=True)
+        error_msg = "Prediction failed. Check logs for details."
+        return error_msg, None, "Error", None, None
+    
+def format_summary(original: str, summary:str) ->str:
+    """Format original and summary text for display"""
+    html = f"""
+    <div style="padding: 10px; border-radius: 5px;">
+        <h3>Original Text</h3>
+        <p style="background-color: #f0f0f0; padding: 10px; border-radius: 3px;">
+            {original}
+        </p>
+        <h3>Summary</h3>
+        <p style="background-color: #e6f3ff; padding: 10px; border-radius: 3px;">
+            {summary}
+        </p>
+    </div>
+    """
+    return html
+
+def create_emotion_plot(emotions: EmotionPrediction | dict[str, Sequence[float] | Sequence[str]]) -> Figure | None:
+    """
+    Create bar plot for emotion predictions.
+    Args:
+        emotions: Dict with 'labels' and 'scores' keys
+    """
+    if isinstance(emotions, EmotionPrediction):
+        labels = emotions.labels
+        scores = emotions.scores
+    else:
+        labels = list(emotions.get("labels", []))
+        scores = list(emotions.get("scores", []))
+
+    if not labels or not scores:
+        return None
+
+    df = pd.DataFrame({
+        "Emotion": labels,
+        "Probability": scores,
+    })
+    fig, ax = plt.subplots(figsize=(8, 5))
+    colors = sns.color_palette("Set2", len(labels))
+    bars = ax.barh(df["Emotion"], df["Probability"], color=colors)
+    ax.set_xlabel("Probability", fontsize=12)
+    ax.set_ylabel("Emotion", fontsize=12)
+    ax.set_title("Emotion Detection Results", fontsize=14, fontweight="bold")
+    ax.set_xlim(0, 1)
+    for bar in bars:
+        width = bar.get_width()
+        ax.text(
+            width,
+            bar.get_y() + bar.get_height() / 2,
+            f"{width:.2%}",
+            ha="left",
+            va="center",
+            fontsize=10,
+            bbox=dict(boxstyle="round,pad=0.3", facecolor="white", alpha=0.7),
+        )
+    plt.tight_layout()
+    return fig
+
+def format_topic(topic: TopicPrediction | dict[str, float | str]) -> str:
+    """
+    Format topic prediction output.
+    
+    Args:
+        topic: Dict with 'label' and 'score' keys
+    """
+    if isinstance(topic, TopicPrediction):
+        label = topic.label
+        score = topic.confidence
+    else:
+        label = str(topic.get("label", "Unknown"))
+        score = float(topic.get("score", 0.0))
+    output = f"""
+    ### Predicted Topic
+        
+    **{label}**
+        
+    Confidence: {score:.2%}
+    """
+    return output
+
+def _clean_tokens(tokens: Iterable[str]) -> list[str]:
+    cleaned: list[str] = []
+    for token in tokens:
+        item = token.replace("Ġ", " ").replace("▁", " ")
+        cleaned.append(item.strip() if item.strip() else token)
+    return cleaned
+
+
+def create_attention_heatmap(text: str, summary: str, pipeline: InferencePipeline) -> Figure | None:
+    """Generate a seaborn heatmap of decoder cross-attention averaged over heads."""
+    if not summary:
+        return None
+    try:
+        batch = pipeline.preprocessor.batch_encode([text])
+        batch = pipeline._batch_to_device(batch)
+        src_ids = batch.input_ids
+        src_mask = batch.attention_mask
+        encoder_mask = None
+        if src_mask is not None:
+            encoder_mask = src_mask.unsqueeze(1) & src_mask.unsqueeze(2)
+
+        with torch.inference_mode():
+            memory = pipeline.model.encoder(src_ids, mask=encoder_mask)
+            target_enc = pipeline.tokenizer.batch_encode([summary])
+            target_ids = target_enc["input_ids"].to(pipeline.device)
+            target_mask = target_enc["attention_mask"].to(pipeline.device)
+            target_len = int(target_mask.sum().item())
+            decoder_inputs = pipeline.tokenizer.prepare_decoder_inputs(target_ids)
+            decoder_inputs = decoder_inputs[:, :target_len].to(pipeline.device)
+            target_ids = target_ids[:, :target_len]
+            memory_mask = src_mask.to(pipeline.device) if src_mask is not None else None
+            _, attn_list = pipeline.model.decoder(
+                decoder_inputs,
+                memory,
+                memory_mask=memory_mask,
+                collect_attn=True,
+            )
+        if not attn_list:
+            return None
+        cross_attn = attn_list[-1]["cross"]  # (B, heads, T, S)
+        attn_matrix = cross_attn.mean(dim=1)[0].detach().cpu().numpy()
+
+        source_len = batch.lengths[0]
+        attn_matrix = attn_matrix[:target_len, :source_len]
+
+        source_ids = src_ids[0, :source_len].tolist()
+        target_id_list = target_ids[0].tolist()
+
+        special_ids = {
+            pipeline.tokenizer.pad_token_id,
+            pipeline.tokenizer.bos_token_id,
+            pipeline.tokenizer.eos_token_id,
+        }
+        keep_indices = [index for index, token_id in enumerate(target_id_list) if token_id not in special_ids]
+        if not keep_indices:
+            return None
+
+        pruned_matrix = attn_matrix[keep_indices, :]
+        tokenizer_impl = pipeline.tokenizer.tokenizer
+        convert_tokens = getattr(tokenizer_impl, "convert_ids_to_tokens", None)
+        if convert_tokens is None:
+            logger.warning("Tokenizer does not expose convert_ids_to_tokens; skipping attention heatmap.")
+            return None
+
+        summary_tokens_raw = convert_tokens([target_id_list[index] for index in keep_indices])
+        source_tokens_raw = convert_tokens(source_ids)
+
+        summary_tokens = _clean_tokens(summary_tokens_raw)
+        source_tokens = _clean_tokens(source_tokens_raw)
+
+        height = max(4.0, 0.4 * len(summary_tokens))
+        width = max(6.0, 0.4 * len(source_tokens))
+        fig, ax = plt.subplots(figsize=(width, height))
+        sns.heatmap(
+            pruned_matrix,
+            cmap="mako",
+            xticklabels=source_tokens,
+            yticklabels=summary_tokens,
+            ax=ax,
+            cbar_kws={"label": "Attention"},
+        )
+        ax.set_xlabel("Input Tokens")
+        ax.set_ylabel("Summary Tokens")
+        ax.set_title("Cross-Attention (decoder last layer)")
+        ax.tick_params(axis="x", rotation=90)
+        ax.tick_params(axis="y", rotation=0)
+        fig.tight_layout()
+        return fig
+
+    except Exception as exc:
+        logger.error("Unable to build attention heatmap: %s", exc, exc_info=True)
+        return None
+
+
+def prepare_download(
+    text: str,
+    summary: str,
+    emotions: EmotionPrediction | dict[str, Sequence[float] | Sequence[str]],
+    topic: TopicPrediction | dict[str, float | str],
+) -> str:
+    """Prepare JSON data for download."""
+    if isinstance(emotions, EmotionPrediction):
+        emotion_payload = {
+            "labels": list(emotions.labels),
+            "scores": list(emotions.scores),
+        }
+    else:
+        emotion_payload = emotions
+
+    if isinstance(topic, TopicPrediction):
+        topic_payload = {
+            "label": topic.label,
+            "confidence": topic.confidence,
+        }
+    else:
+        topic_payload = topic
+
+    data = {
+        "original_text": text,
+        "summary": summary,
+        "emotions": emotion_payload,
+        "topic": topic_payload,
+    }
+    return json.dumps(data, indent=2)
+
+# Sample data for the demo
+SAMPLE_TEXT = """
+Artificial intelligence is rapidly transforming the technology landscape. 
+Machine learning algorithms are now capable of processing vast amounts of data, 
+identifying patterns, and making predictions with unprecedented accuracy. 
+From healthcare diagnostics to financial forecasting, AI applications are 
+revolutionizing industries worldwide. However, ethical considerations around 
+privacy, bias, and transparency remain critical challenges that must be addressed 
+as these technologies continue to evolve.
+"""
+
+def create_interface():
+    with gr.Blocks(title="LexiMind Demo", theme=Soft()) as demo:
+        gr.Markdown("""
+        # LexiMind NLP Pipeline Demo
+        
+        **Full pipleine for text summarization, emotion detection, and topic prediction.**
+        
+        Enter text below and adjust compressoin to see the results.
+        """)
+        with gr.Row():
+            # Left column - Input
+            with gr.Column(scale=1):
+                gr.Markdown("### Input")
+                input_text = gr.Textbox(
+                    label="Enter text",
+                    placeholder="Paste or type your text here...",
+                    lines=10,
+                    value=SAMPLE_TEXT
+                )
+                token_count = gr.Textbox(
+                    label="Token Count",
+                    value="Tokens: 0",
+                    interactive=False
+                )
+                compression = gr.Slider(
+                    minimum=20,
+                    maximum=80,
+                    value=50,
+                    step=5,
+                    label="Compression %",
+                    info="Higher = shorter summary"
+                )
+                predict_btn = gr.Button("🚀 Analyze", variant="primary", size="lg")
+            # Right column - Outputs
+            with gr.Column(scale=2):
+                gr.Markdown("### Result")
+                with gr.Tabs():
+                    with gr.TabItem("Summary"):
+                        summary_output = gr.HTML(label="Summary")
+                    with gr.TabItem("Emotions"):
+                        emotion_output = gr.Plot(label="Emotion Analysis")
+                    with gr.TabItem("Topic"):
+                        topic_output = gr.Markdown(label="Topic Prediction")
+                    with gr.TabItem("Attention Heatmap"):
+                        attention_output = gr.Plot(label="Attention Weights")
+                        gr.Markdown("*Visualizes which parts of the input the model focused on.*")
+                # Download section
+                gr.Markdown("### Export Results")
+                download_data = gr.Textbox(visible=False)
+                download_btn = gr.DownloadButton(
+                    "Download Results (JSON)",
+                    visible=True
+                )
+            # Event Handlers
+            input_text.change(
+                fn=count_tokens,
+                inputs=[input_text],
+                outputs=[token_count]
+            )
+            predict_btn.click(
+                fn=predict,
+                inputs=[input_text, compression],
+                outputs=[summary_output, emotion_output, topic_output, attention_output, download_data]
+            ).then(
+                fn=lambda x: gr.DownloadButton("Download Results (JSON)", value=x, visible=True),
+                inputs=[download_data],
+                outputs=[download_btn]
+            )
+            # Examples
+            gr.Examples(
+            examples=[
+                [SAMPLE_TEXT, 50],
+                ["Climate change poses significant risks to global ecosystems. Rising temperatures, melting ice caps, and extreme weather events are becoming more frequent. Scientists urge immediate action to reduce carbon emissions and transition to renewable energy sources.", 40],
+                ],
+                inputs=[input_text, compression],
+                label="Try these examples:"
+            )
+        return demo
+    
+    if __name__ == "__main__":
+        try:
+            # Pre-load pipeline
+            get_pipeline()
+        
+            # Create and launch interface
+            demo = create_interface()
+            demo.queue()  # Enable queuing for better responsiveness
+            demo.launch(
+            share=True,
+            server_name="0.0.0.0",
+            server_port=7860,
+            show_error=True
+            )
+        
+        except Exception as e:
+            logger.error(f"Failed to launch demo: {e}", exc_info=True)
+            print(f"Error: {e}")
+            sys.exit(1)
+        
+