app.py

import gradio as gr
from transformers import pipeline
import time

# ==========================================
# MODEL CONFIGURATION
# ==========================================
MODEL_NAME = "ENTUM-AI/FinBERT-Pro"

print(f"Loading model: {MODEL_NAME}...")
try:
    classifier = pipeline("text-classification", model=MODEL_NAME, top_k=3)
    print("Model loaded successfully!")
except Exception as e:
    print(f"Error loading model: {e}")
    classifier = None


# ==========================================
# PREDICTION LOGIC
# ==========================================
SENTIMENT_CONFIG = {
    "Positive": {"color": "#16a34a", "bg": "#f0fdf4", "icon": "📈", "bar": "#22c55e"},
    "Negative": {"color": "#dc2626", "bg": "#fef2f2", "icon": "📉", "bar": "#ef4444"},
    "Neutral":  {"color": "#2563eb", "bg": "#eff6ff", "icon": "➖", "bar": "#3b82f6"},
}


def predict_single(text):
    """Classify a single financial text."""
    if not text or not text.strip():
        return create_empty_result()

    if classifier is None:
        return create_error_result()

    start = time.time()
    results = classifier(text.strip())[0]
    elapsed = (time.time() - start) * 1000

    top = results[0]
    return create_result_html(text.strip(), top["label"], results, elapsed)


def predict_batch(texts):
    """Classify multiple financial texts (one per line)."""
    if not texts or not texts.strip():
        return "<p style='color:#94a3b8; text-align:center;'>Enter financial texts, one per line.</p>"

    if classifier is None:
        return create_error_result()

    lines = [line.strip() for line in texts.strip().split("\n") if line.strip()]
    if not lines:
        return "<p style='color:#94a3b8; text-align:center;'>No valid texts found.</p>"

    start = time.time()
    all_results = classifier(lines)
    elapsed = (time.time() - start) * 1000

    counts = {"Positive": 0, "Negative": 0, "Neutral": 0}
    html_parts = []

    for text, results in zip(lines, all_results):
        top = results[0]
        label = top["label"]
        score = top["score"]
        counts[label] = counts.get(label, 0) + 1

        cfg = SENTIMENT_CONFIG.get(label, SENTIMENT_CONFIG["Neutral"])
        bar_width = int(score * 100)

        html_parts.append(f"""
        <div style="
            background: {cfg['bg']};
            border: 1px solid {cfg['color']}22;
            border-left: 4px solid {cfg['color']};
            border-radius: 12px;
            padding: 16px 20px;
            margin-bottom: 10px;
        ">
            <div style="display:flex; justify-content:space-between; align-items:center; margin-bottom:8px;">
                <span style="color:#1e293b; font-size:14px; flex:1; margin-right:12px;">{cfg['icon']} {text}</span>
                <span style="
                    background: {cfg['color']}15;
                    color: {cfg['color']};
                    padding: 4px 12px;
                    border-radius: 20px;
                    font-size: 12px;
                    font-weight: 700;
                    letter-spacing: 0.5px;
                    white-space: nowrap;
                ">{label.upper()} {score:.0%}</span>
            </div>
            <div style="background:#e2e8f0; border-radius:6px; height:6px; overflow:hidden;">
                <div style="width:{bar_width}%; height:100%; background:linear-gradient(90deg, {cfg['bar']}aa, {cfg['bar']}); border-radius:6px;"></div>
            </div>
        </div>
        """)

    # Summary card
    total = len(lines)
    summary = f"""
    <div style="
        background: #ffffff;
        border: 1px solid #e2e8f0;
        border-radius: 14px;
        padding: 20px 24px;
        margin-bottom: 16px;
        text-align: center;
        box-shadow: 0 1px 3px rgba(0,0,0,0.06);
    ">
        <span style="color:#64748b; font-size:12px; text-transform:uppercase; letter-spacing:1px;">Batch Sentiment Analysis</span>
        <div style="display:flex; justify-content:center; gap:24px; margin-top:12px;">
            <div>
                <div style="color:#16a34a; font-size:24px; font-weight:800;">📈 {counts.get('Positive', 0)}</div>
                <div style="color:#64748b; font-size:12px;">Positive</div>
            </div>
            <div>
                <div style="color:#2563eb; font-size:24px; font-weight:800;">➖ {counts.get('Neutral', 0)}</div>
                <div style="color:#64748b; font-size:12px;">Neutral</div>
            </div>
            <div>
                <div style="color:#dc2626; font-size:24px; font-weight:800;">📉 {counts.get('Negative', 0)}</div>
                <div style="color:#64748b; font-size:12px;">Negative</div>
            </div>
        </div>
        <div style="color:#94a3b8; font-size:12px; margin-top:10px;">{total} texts analyzed in {elapsed:.0f}ms</div>
    </div>
    """

    return summary + "\n".join(html_parts)


# ==========================================
# HTML RESULT BUILDERS
# ==========================================
def create_result_html(text, top_label, all_scores, elapsed_ms):
    cfg = SENTIMENT_CONFIG.get(top_label, SENTIMENT_CONFIG["Neutral"])

    # Build sentiment bars for all 3 classes
    bars_html = ""
    for item in all_scores:
        lbl = item["label"]
        sc = item["score"]
        c = SENTIMENT_CONFIG.get(lbl, SENTIMENT_CONFIG["Neutral"])
        pct = int(sc * 100)
        bars_html += f"""
        <div style="margin-bottom:10px;">
            <div style="display:flex; justify-content:space-between; margin-bottom:4px;">
                <span style="color:#475569; font-size:13px; font-weight:600;">{c['icon']} {lbl}</span>
                <span style="color:{c['color']}; font-weight:700; font-size:14px;">{sc:.1%}</span>
            </div>
            <div style="background:#f1f5f9; border-radius:6px; height:8px; overflow:hidden;">
                <div style="width:{pct}%; height:100%; background:linear-gradient(90deg, {c['bar']}aa, {c['bar']}); border-radius:6px;"></div>
            </div>
        </div>
        """

    # Gradient background based on sentiment
    if top_label == "Positive":
        gradient = "linear-gradient(135deg, #dcfce7, #bbf7d0, #86efac)"
        text_color = "#166534"
    elif top_label == "Negative":
        gradient = "linear-gradient(135deg, #fee2e2, #fecaca, #fca5a5)"
        text_color = "#991b1b"
    else:
        gradient = "linear-gradient(135deg, #dbeafe, #bfdbfe, #93c5fd)"
        text_color = "#1e40af"

    top_score = all_scores[0]["score"]

    return f"""
    <div style="font-family: 'Inter', 'Segoe UI', sans-serif;">
        <div style="
            background: {gradient};
            border-radius: 20px;
            padding: 32px;
            text-align: center;
            margin-bottom: 20px;
            border: 1px solid {cfg['color']}22;
            box-shadow: 0 4px 24px {cfg['color']}15;
        ">
            <div style="font-size: 48px; margin-bottom: 8px;">{cfg['icon']}</div>
            <div style="
                color: {text_color};
                font-size: 22px;
                font-weight: 800;
                letter-spacing: 2px;
                margin-bottom: 6px;
            ">{top_label.upper()} SENTIMENT</div>
            <div style="color: {text_color}99; font-size: 14px;">Confidence: {top_score:.1%}</div>
        </div>

        <div style="
            background: #ffffff;
            border: 1px solid #e2e8f0;
            border-radius: 16px;
            padding: 24px;
            box-shadow: 0 1px 3px rgba(0,0,0,0.06);
        ">
            <div style="margin-bottom: 20px;">
                <span style="color: #64748b; font-size: 11px; text-transform: uppercase; letter-spacing: 1px;">Analyzed Text</span>
                <div style="color: #1e293b; font-size: 15px; margin-top: 6px; font-style: italic;">"{text}"</div>
            </div>

            <div style="margin-bottom: 16px;">
                <span style="color: #64748b; font-size: 11px; text-transform: uppercase; letter-spacing: 1px; margin-bottom: 10px; display:block;">Score Breakdown</span>
                {bars_html}
            </div>

            <div style="
                display: flex;
                justify-content: center;
                gap: 24px;
                padding-top: 12px;
                border-top: 1px solid #f1f5f9;
            ">
                <div style="text-align: center;">
                    <span style="color: #94a3b8; font-size: 11px; text-transform: uppercase; letter-spacing: 0.5px;">Model</span>
                    <div style="color: #6366f1; font-size: 13px; font-weight: 600; margin-top: 2px;">FinBERT-Pro</div>
                </div>
                <div style="text-align: center;">
                    <span style="color: #94a3b8; font-size: 11px; text-transform: uppercase; letter-spacing: 0.5px;">Latency</span>
                    <div style="color: #0891b2; font-size: 13px; font-weight: 600; margin-top: 2px;">{elapsed_ms:.0f}ms</div>
                </div>
                <div style="text-align: center;">
                    <span style="color: #94a3b8; font-size: 11px; text-transform: uppercase; letter-spacing: 0.5px;">Classes</span>
                    <div style="color: #d97706; font-size: 13px; font-weight: 600; margin-top: 2px;">3</div>
                </div>
            </div>
        </div>
    </div>
    """


def create_empty_result():
    return """
    <div style="
        text-align: center;
        padding: 60px 24px;
        color: #94a3b8;
    ">
        <div style="font-size: 48px; margin-bottom: 12px;">💹</div>
        <div style="font-size: 16px; font-weight: 600; color: #475569;">Awaiting Input</div>
        <div style="font-size: 13px; margin-top: 4px;">Enter a financial text above and click <b>Analyze</b></div>
    </div>
    """


def create_error_result():
    return """
    <div style="
        text-align: center;
        padding: 40px 24px;
        background: #fef2f2;
        border-radius: 16px;
        border: 1px solid #fecaca;
    ">
        <div style="font-size: 36px; margin-bottom: 8px;">⚠️</div>
        <div style="color: #dc2626; font-size: 15px; font-weight: 600;">Model Not Available</div>
        <div style="color: #64748b; font-size: 13px; margin-top: 4px;">Please wait while the model loads or try refreshing.</div>
    </div>
    """


# ==========================================
# CUSTOM CSS
# ==========================================
CUSTOM_CSS = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700;800&display=swap');

* { font-family: 'Inter', 'Segoe UI', sans-serif !important; }

.gradio-container {
    max-width: 960px !important;
    margin: 0 auto !important;
    background: linear-gradient(180deg, #f8fafc 0%, #f1f5f9 50%, #e2e8f0 100%) !important;
}

.main-header {
    text-align: center;
    padding: 40px 20px 20px;
}

.main-header h1 {
    background: linear-gradient(135deg, #059669, #0891b2, #2563eb);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    font-size: 2.5rem !important;
    font-weight: 800 !important;
    margin-bottom: 8px !important;
    letter-spacing: -0.5px;
}

.main-header p {
    color: #64748b !important;
    font-size: 15px !important;
}

.model-badge {
    display: inline-block;
    background: linear-gradient(135deg, #ecfdf5, #e0f2fe);
    border: 1px solid #a7f3d0;
    color: #047857 !important;
    padding: 6px 16px;
    border-radius: 24px;
    font-size: 13px !important;
    font-weight: 600;
    letter-spacing: 0.5px;
    margin-top: 12px;
}

footer { display: none !important; }

.tab-nav button {
    color: #64748b !important;
    font-weight: 600 !important;
    font-size: 14px !important;
}

.tab-nav button.selected {
    color: #059669 !important;
    border-color: #059669 !important;
}
"""


# ==========================================
# GRADIO UI
# ==========================================
with gr.Blocks(
    css=CUSTOM_CSS,
    title="FinBERT-Pro — Financial Sentiment Analyzer",
    theme=gr.themes.Soft(
        primary_hue="emerald",
        secondary_hue="cyan",
        neutral_hue="slate",
    ),
) as demo:

    # Header
    gr.HTML("""
    <div class="main-header">
        <h1>💹 FinBERT-Pro</h1>
        <p>Financial sentiment analysis powered by <b>FinBERT</b>, fine-tuned on 3 expert-annotated datasets</p>
        <span class="model-badge">🧠 ENTUM-AI / FinBERT-Pro</span>
    </div>
    """)

    with gr.Tabs():
        # --- Tab 1: Single Analysis ---
        with gr.Tab("🔍 Single Analysis"):
            with gr.Row():
                with gr.Column(scale=3):
                    single_input = gr.Textbox(
                        label="Financial Text",
                        placeholder="e.g. Stock price soars on record-breaking earnings report",
                        lines=2,
                        max_lines=4,
                    )
                    single_btn = gr.Button("⚡ Analyze Sentiment", variant="primary", size="lg")
                with gr.Column(scale=4):
                    single_output = gr.HTML(value=create_empty_result())

            gr.Examples(
                examples=[
                    ["Stock price soars on record-breaking earnings report"],
                    ["Revenue decline signals weakening market position"],
                    ["Company announces quarterly earnings results"],
                    ["Shares surge 15% after strong Q3 revenue growth"],
                    ["Major layoffs expected as company restructures operations"],
                    ["The board of directors met to discuss routine operations"],
                    ["Bankruptcy filing raises concerns about long-term viability"],
                    ["Profit margins improved significantly driven by cost optimization"],
                ],
                inputs=single_input,
                label="📋 Try these examples",
            )

            single_btn.click(fn=predict_single, inputs=single_input, outputs=single_output)
            single_input.submit(fn=predict_single, inputs=single_input, outputs=single_output)

        # --- Tab 2: Batch Analysis ---
        with gr.Tab("📊 Batch Analysis"):
            gr.Markdown("Paste multiple financial texts — **one per line** — for batch sentiment classification.")
            with gr.Row():
                with gr.Column(scale=2):
                    batch_input = gr.Textbox(
                        label="Financial Texts (one per line)",
                        placeholder="Headline 1\nHeadline 2\nHeadline 3",
                        lines=8,
                        max_lines=20,
                    )
                    batch_btn = gr.Button("⚡ Analyze All", variant="primary", size="lg")
                with gr.Column(scale=3):
                    batch_output = gr.HTML(
                        value="<p style='color:#94a3b8; text-align:center; padding:40px;'>Results will appear here.</p>"
                    )

            batch_btn.click(fn=predict_batch, inputs=batch_input, outputs=batch_output)

        # --- Tab 3: About ---
        with gr.Tab("ℹ️ About"):
            gr.HTML("""
            <div style="
                background: #ffffff;
                border: 1px solid #e2e8f0;
                border-radius: 20px;
                padding: 36px;
                color: #1e293b;
                box-shadow: 0 1px 3px rgba(0,0,0,0.06);
            ">
                <h2 style="
                    background: linear-gradient(135deg, #059669, #0891b2);
                    -webkit-background-clip: text;
                    -webkit-text-fill-color: transparent;
                    font-size: 24px;
                    margin-bottom: 24px;
                ">About FinBERT-Pro</h2>

                <p style="color:#475569; font-size:14px; line-height:1.7; margin-bottom:20px;">
                    An improved financial sentiment model built on
                    <a href="https://huggingface.co/ProsusAI/finbert" style="color:#059669; text-decoration:none; font-weight:600;">ProsusAI/FinBERT</a>.
                    Fine-tuned on 3 expert-annotated financial datasets for more robust sentiment classification.
                </p>

                <table style="width:100%; border-collapse:separate; border-spacing:0 8px;">
                    <tr>
                        <td style="color:#64748b; padding:8px 16px; font-size:13px; width:35%;">Base Model</td>
                        <td style="color:#1e293b; padding:8px 16px; font-size:14px; font-weight:600;">ProsusAI/FinBERT (BERT-based)</td>
                    </tr>
                    <tr>
                        <td style="color:#64748b; padding:8px 16px; font-size:13px;">Task</td>
                        <td style="color:#1e293b; padding:8px 16px; font-size:14px; font-weight:600;">3-Class Sentiment (Positive / Negative / Neutral)</td>
                    </tr>
                    <tr>
                        <td style="color:#64748b; padding:8px 16px; font-size:13px;">Language</td>
                        <td style="color:#1e293b; padding:8px 16px; font-size:14px; font-weight:600;">English</td>
                    </tr>
                    <tr>
                        <td style="color:#64748b; padding:8px 16px; font-size:13px;">Max Input</td>
                        <td style="color:#1e293b; padding:8px 16px; font-size:14px; font-weight:600;">128 tokens</td>
                    </tr>
                    <tr>
                        <td style="color:#64748b; padding:8px 16px; font-size:13px;">License</td>
                        <td style="color:#1e293b; padding:8px 16px; font-size:14px; font-weight:600;">Apache 2.0</td>
                    </tr>
                </table>

                <h3 style="color:#059669; margin-top:28px; margin-bottom:12px; font-size:16px;">📊 Training Data</h3>
                <p style="color:#475569; font-size:13px; line-height:1.7;">
                    Fine-tuned on 3 expert-annotated datasets (~14.6K total samples):
                    <b style="color:#047857;">FinanceInc/auditor_sentiment</b>,
                    <b style="color:#047857;">nickmuchi/financial-classification</b>, and
                    <b style="color:#047857;">warwickai/financial_phrasebank_mirror</b>.
                </p>

                <h3 style="color:#059669; margin-top:28px; margin-bottom:12px; font-size:16px;">🔍 What's Different from FinBERT?</h3>
                <div style="display:grid; grid-template-columns:1fr 1fr 1fr; gap:10px; margin-top:12px;">
                    <div style="background:#f8fafc; border:1px solid #e2e8f0; padding:14px; border-radius:10px; font-size:13px; text-align:center;">
                        <span style="font-size:20px;">📚</span><br>
                        <b style="color:#1e293b;">Multi-Source</b><br>
                        <span style="color:#64748b;">3 datasets vs 1</span>
                    </div>
                    <div style="background:#f8fafc; border:1px solid #e2e8f0; padding:14px; border-radius:10px; font-size:13px; text-align:center;">
                        <span style="font-size:20px;">⚖️</span><br>
                        <b style="color:#1e293b;">Class-Weighted</b><br>
                        <span style="color:#64748b;">Handles imbalance</span>
                    </div>
                    <div style="background:#f8fafc; border:1px solid #e2e8f0; padding:14px; border-radius:10px; font-size:13px; text-align:center;">
                        <span style="font-size:20px;">🎯</span><br>
                        <b style="color:#1e293b;">Robust</b><br>
                        <span style="color:#64748b;">Better generalization</span>
                    </div>
                </div>

                <h3 style="color:#059669; margin-top:28px; margin-bottom:12px; font-size:16px;">🐍 Python API</h3>
                <pre style="
                    background:#f8fafc;
                    border:1px solid #e2e8f0;
                    border-radius:12px;
                    padding:20px;
                    color:#1e293b;
                    font-size:13px;
                    overflow-x:auto;
                    font-family: 'Fira Code', 'Cascadia Code', monospace !important;
                "><span style="color:#059669">from</span> transformers <span style="color:#059669">import</span> pipeline

classifier = pipeline(<span style="color:#d97706">"text-classification"</span>,
                       model=<span style="color:#d97706">"ENTUM-AI/FinBERT-Pro"</span>)

result = classifier(<span style="color:#d97706">"Stock price soars on earnings"</span>)
<span style="color:#94a3b8"># [{'label': 'Positive', 'score': 0.99}]</span></pre>
            </div>
            """)


# Launch
demo.launch()