app.py

import gradio as gr
from llama_cpp import Llama
import re

# ==========================================
# 1. OPTIMIZED MODEL LOADING
# ==========================================
print("Loading model...")
llm = Llama(
    model_path="qwen_rewrite_q4km.gguf",
    n_ctx=4096,           # Increased context window for safer processing of larger drafts
    n_threads=2,          
    n_batch=512,          # Optimized batch size for prompt processing speed
    verbose=False
)
print("Model loaded successfully!")

# ==========================================
# 2. ADVANCED INFERENCE & PROMPT PIPELINE
# ==========================================
def rewrite_text(instruction, original_text):
    """
    Upgraded Rewrite Engine applying principles from:
    - "Principled Instructions Are All You Need" (Bsharat et al., 2023)
    - "The Curious Case of Neural Text Degeneration" (Holtzman et al., 2019)
    """
    
    # STRUCTURAL PROMPTING: Using Markdown headers creates strict logical 
    # boundaries for small models, preventing context confusion.
    prompt = f"""<|im_start|>system
You are an elite, highly-paid copywriter and text revision engine. Your sole function is to rewrite text based on specific instructions.
OPERATING RULES:
1. Apply the user's instruction precisely and completely.
2. Enhance vocabulary, flow, and clarity automatically.
3. NEVER output conversational filler, pleasantries, or meta-commentary.
4. Output STRICTLY the finalized text and absolutely nothing else.
<|im_end|>
<|im_start|>user
# INSTRUCTION
{instruction.strip()}

# ORIGINAL TEXT
{original_text.strip()}
<|im_end|>
<|im_start|>assistant
"""

    # ADVANCED DECODING PARAMETERS
    response = llm(
        prompt,
        max_tokens=1024,          # Increased to ensure longer outputs aren't truncated
        stop=["<|im_end|>", "<|im_start|>"], 
        
        # SAMPLING OPTIMIZATIONS
        temperature=0.35,         # Slight bump for vocabulary variance
        top_p=0.85,               # Nucleus Sampling: Truncates the worst 15% of possible words
        min_p=0.05,               # Min-P: Dynamically cuts off words that are less than 5% as likely as the best word (Prevents hallucinations)
        top_k=40,                 # Hard limit on top choices
        
        # PENALTIES (Crucial for 1.5B models to prevent looping/blandness)
        repeat_penalty=1.15,      # Penalizes the model for repeating exactly the same sequence
        presence_penalty=0.1,     # Encourages the model to use new words it hasn't used yet
        frequency_penalty=0.1,    # Discourages repeating the same words frequently
        
        echo=False
    )

    result = response["choices"][0]["text"].strip()

    # PROGRAMMATIC SANITIZATION
    # Small models occasionally ignore the "no conversational filler" rule. 
    # This Regex safely strips out "Sure, here is the text:\n" behavior if it occurs.
    filler_patterns = [
        r"^(here is|here's|sure|certainly|the rewritten|rewritten).*?:",
        r"^(here is|here's|sure|certainly|the rewritten|rewritten).*?\n"
    ]
    for pattern in filler_patterns:
        result = re.sub(pattern, "", result, flags=re.IGNORECASE).strip()

    return result


# ──────────────────────────────────────────────
# Aurora Noir — Ultra-Premium Dark UI (UNTOUCHED)
# ──────────────────────────────────────────────
aurora_css = """
@import url('https://fonts.googleapis.com/css2?family=Space+Grotesk:wght@300;400;500;600;700&family=JetBrains+Mono:wght@400;500&display=swap');

/* ═══ Global Canvas ═══ */
body, .gradio-container {
    font-family: 'Space Grotesk', -apple-system, BlinkMacSystemFont, sans-serif !important;
    background: #09090b !important;
    min-height: 100vh;
    position: relative;
}

/* ═══ Ambient Background Glow ═══ */
body::before {
    content: '';
    position: fixed;
    top: -40%;
    left: -20%;
    width: 80%;
    height: 80%;
    background: radial-gradient(ellipse at center, rgba(120, 80, 255, 0.06) 0%, transparent 70%);
    pointer-events: none;
    z-index: 0;
}

body::after {
    content: '';
    position: fixed;
    bottom: -30%;
    right: -20%;
    width: 70%;
    height: 70%;
    background: radial-gradient(ellipse at center, rgba(236, 72, 153, 0.04) 0%, transparent 70%);
    pointer-events: none;
    z-index: 0;
}

.gradio-container {
    max-width: 1100px !important;
    margin: 0 auto !important;
    padding: 48px 24px !important;
    position: relative;
    z-index: 1;
}

/* ═══ Main Glass Container ═══ */
.main-container {
    background: linear-gradient(165deg, rgba(18, 18, 24, 0.95) 0%, rgba(12, 12, 18, 0.98) 100%) !important;
    border: 1px solid rgba(255, 255, 255, 0.04) !important;
    border-radius: 28px !important;
    box-shadow:
        0 0 0 0.5px rgba(255, 255, 255, 0.03) inset,
        0 40px 120px -20px rgba(0, 0, 0, 0.8),
        0 0 80px rgba(120, 80, 255, 0.03) !important;
    padding: 0 !important;
    overflow: hidden !important;
    position: relative !important;
}

/* Top accent line */
.main-container::before {
    content: '';
    position: absolute;
    top: 0;
    left: 10%;
    right: 10%;
    height: 1px;
    background: linear-gradient(90deg,
        transparent 0%,
        rgba(139, 92, 246, 0.4) 25%,
        rgba(236, 72, 153, 0.4) 50%,
        rgba(59, 130, 246, 0.4) 75%,
        transparent 100%);
    z-index: 10;
}

/* ═══ Top Navigation Bar ═══ */
.top-bar {
    padding: 18px 32px !important;
    display: flex !important;
    align-items: center !important;
    justify-content: space-between !important;
    border-bottom: 1px solid rgba(255, 255, 255, 0.03) !important;
    background: rgba(255, 255, 255, 0.01) !important;
}

.logo-group {
    display: flex !important;
    align-items: center !important;
    gap: 12px !important;
}

.logo-icon {
    width: 36px !important;
    height: 36px !important;
    border-radius: 10px !important;
    background: linear-gradient(135deg, #7c3aed 0%, #ec4899 100%) !important;
    display: flex !important;
    align-items: center !important;
    justify-content: center !important;
    font-size: 16px !important;
    box-shadow: 0 4px 16px rgba(124, 58, 237, 0.3) !important;
}

.logo-text {
    font-size: 15px !important;
    font-weight: 600 !important;
    color: rgba(255, 255, 255, 0.85) !important;
    letter-spacing: -0.02em !important;
}

.status-badge {
    display: flex !important;
    align-items: center !important;
    gap: 6px !important;
    background: rgba(34, 197, 94, 0.08) !important;
    border: 1px solid rgba(34, 197, 94, 0.15) !important;
    padding: 5px 14px !important;
    border-radius: 100px !important;
    font-size: 11px !important;
    font-weight: 500 !important;
    color: rgba(34, 197, 94, 0.8) !important;
    letter-spacing: 0.03em !important;
}

.status-dot {
    width: 6px !important;
    height: 6px !important;
    border-radius: 50% !important;
    background: #22c55e !important;
    box-shadow: 0 0 8px rgba(34, 197, 94, 0.5) !important;
    animation: pulse-dot 2s ease-in-out infinite !important;
}

@keyframes pulse-dot {
    0%, 100% { opacity: 1; box-shadow: 0 0 8px rgba(34, 197, 94, 0.5); }
    50% { opacity: 0.6; box-shadow: 0 0 16px rgba(34, 197, 94, 0.8); }
}

/* ═══ Content Body ═══ */
.content-body {
    padding: 44px 48px 48px !important;
}

/* ═══ Hero Section ═══ */
.hero-block {
    text-align: center !important;
    margin-bottom: 44px !important;
    position: relative !important;
}

.hero-block h1 {
    font-size: 42px !important;
    font-weight: 700 !important;
    color: #fafafa !important;
    margin: 0 0 4px 0 !important;
    letter-spacing: -0.04em !important;
    line-height: 1.15 !important;
}

.hero-block h1 .gradient-word {
    background: linear-gradient(135deg, #a78bfa 0%, #c084fc 25%, #f472b6 50%, #fb923c 75%, #a78bfa 100%) !important;
    background-size: 200% auto !important;
    -webkit-background-clip: text !important;
    -webkit-text-fill-color: transparent !important;
    background-clip: text !important;
    animation: shimmer 4s linear infinite !important;
}

@keyframes shimmer {
    0% { background-position: 0% center; }
    100% { background-position: 200% center; }
}

.hero-block .tagline {
    font-size: 15px !important;
    color: rgba(255, 255, 255, 0.3) !important;
    font-weight: 400 !important;
    line-height: 1.7 !important;
    max-width: 500px !important;
    margin: 16px auto 0 !important;
}

.feature-chips {
    display: flex !important;
    justify-content: center !important;
    gap: 8px !important;
    margin-top: 24px !important;
    flex-wrap: wrap !important;
}

.feature-chip {
    display: inline-flex !important;
    align-items: center !important;
    gap: 5px !important;
    background: rgba(255, 255, 255, 0.03) !important;
    border: 1px solid rgba(255, 255, 255, 0.05) !important;
    padding: 6px 14px !important;
    border-radius: 8px !important;
    font-size: 12px !important;
    color: rgba(255, 255, 255, 0.4) !important;
    font-weight: 500 !important;
    transition: all 0.3s ease !important;
}

.feature-chip:hover {
    border-color: rgba(139, 92, 246, 0.2) !important;
    color: rgba(255, 255, 255, 0.6) !important;
    background: rgba(139, 92, 246, 0.05) !important;
}

.divider-line {
    width: 60px !important;
    height: 1px !important;
    background: linear-gradient(90deg, transparent, rgba(139, 92, 246, 0.3), transparent) !important;
    margin: 32px auto 0 !important;
}

/* ═══ Section Titles ═══ */
.panel-label {
    display: flex !important;
    align-items: center !important;
    gap: 8px !important;
    margin-bottom: 16px !important;
}

.panel-label-icon {
    width: 20px !important;
    height: 20px !important;
    border-radius: 6px !important;
    display: flex !important;
    align-items: center !important;
    justify-content: center !important;
    font-size: 10px !important;
}

.panel-label-icon.input-icon {
    background: rgba(59, 130, 246, 0.1) !important;
    border: 1px solid rgba(59, 130, 246, 0.15) !important;
}

.panel-label-icon.output-icon {
    background: rgba(34, 197, 94, 0.1) !important;
    border: 1px solid rgba(34, 197, 94, 0.15) !important;
}

.panel-label-text {
    font-size: 11px !important;
    font-weight: 600 !important;
    text-transform: uppercase !important;
    letter-spacing: 0.12em !important;
    color: rgba(255, 255, 255, 0.3) !important;
}

/* ═══ Input Cards ═══ */
.glass-card {
    background: rgba(255, 255, 255, 0.015) !important;
    border: 1px solid rgba(255, 255, 255, 0.04) !important;
    border-radius: 16px !important;
    padding: 20px 22px !important;
    margin-bottom: 12px !important;
    transition: all 0.4s cubic-bezier(0.4, 0, 0.2, 1) !important;
    position: relative !important;
    overflow: hidden !important;
}

.glass-card::before {
    content: '';
    position: absolute;
    top: 0;
    left: 0;
    right: 0;
    bottom: 0;
    background: linear-gradient(135deg, rgba(139, 92, 246, 0.02) 0%, transparent 50%);
    opacity: 0;
    transition: opacity 0.4s ease;
    pointer-events: none;
}

.glass-card:hover {
    border-color: rgba(139, 92, 246, 0.12) !important;
    box-shadow: 0 8px 40px rgba(139, 92, 246, 0.04) !important;
}

.glass-card:hover::before {
    opacity: 1;
}

/* ═══ Textboxes ═══ */
.gradio-container textarea,
.gradio-container input[type="text"] {
    background: rgba(0, 0, 0, 0.3) !important;
    border: 1px solid rgba(255, 255, 255, 0.05) !important;
    border-radius: 12px !important;
    color: #e4e4e7 !important;
    font-family: 'Space Grotesk', sans-serif !important;
    font-size: 14px !important;
    padding: 14px 18px !important;
    transition: all 0.3s cubic-bezier(0.4, 0, 0.2, 1) !important;
    line-height: 1.65 !important;
    letter-spacing: -0.005em !important;
}

.gradio-container textarea:focus,
.gradio-container input[type="text"]:focus {
    border-color: rgba(139, 92, 246, 0.35) !important;
    box-shadow:
        0 0 0 3px rgba(139, 92, 246, 0.08),
        0 0 30px rgba(139, 92, 246, 0.06) !important;
    outline: none !important;
    background: rgba(0, 0, 0, 0.4) !important;
}

.gradio-container textarea::placeholder,
.gradio-container input[type="text"]::placeholder {
    color: rgba(255, 255, 255, 0.15) !important;
    font-weight: 400 !important;
}

/* ═══ Labels ═══ */
.gradio-container label,
.gradio-container .label-wrap span {
    font-size: 12.5px !important;
    font-weight: 500 !important;
    color: rgba(255, 255, 255, 0.55) !important;
    letter-spacing: 0.01em !important;
}

/* ═══ Primary Button — Glow Effect ═══ */
.gr-button-primary,
button.primary {
    background: linear-gradient(135deg, #7c3aed 0%, #a855f7 50%, #7c3aed 100%) !important;
    background-size: 200% auto !important;
    border: none !important;
    border-radius: 14px !important;
    color: white !important;
    font-family: 'Space Grotesk', sans-serif !important;
    font-size: 14.5px !important;
    font-weight: 600 !important;
    padding: 14px 32px !important;
    cursor: pointer !important;
    transition: all 0.4s cubic-bezier(0.4, 0, 0.2, 1) !important;
    box-shadow:
        0 6px 24px rgba(124, 58, 237, 0.25),
        0 0 0 0.5px rgba(255, 255, 255, 0.06) inset !important;
    letter-spacing: -0.005em !important;
    width: 100% !important;
    margin-top: 6px !important;
    position: relative !important;
    overflow: hidden !important;
}

.gr-button-primary::before,
button.primary::before {
    content: '';
    position: absolute;
    top: 0;
    left: -100%;
    width: 100%;
    height: 100%;
    background: linear-gradient(90deg, transparent, rgba(255,255,255,0.08), transparent);
    transition: left 0.5s ease;
}

.gr-button-primary:hover,
button.primary:hover {
    background-position: right center !important;
    box-shadow:
        0 12px 40px rgba(124, 58, 237, 0.35),
        0 0 60px rgba(124, 58, 237, 0.1),
        0 0 0 0.5px rgba(255, 255, 255, 0.1) inset !important;
    transform: translateY(-2px) !important;
}

.gr-button-primary:hover::before,
button.primary:hover::before {
    left: 100%;
}

.gr-button-primary:active,
button.primary:active {
    transform: translateY(0px) scale(0.99) !important;
    box-shadow: 0 4px 16px rgba(124, 58, 237, 0.25) !important;
}

/* ═══ Output Panel ═══ */
.output-panel {
    background: rgba(255, 255, 255, 0.015) !important;
    border: 1px solid rgba(255, 255, 255, 0.04) !important;
    border-radius: 16px !important;
    padding: 22px !important;
    position: relative !important;
    min-height: 200px !important;
    overflow: hidden !important;
}

.output-panel::before {
    content: '';
    position: absolute;
    top: 0;
    left: 15%;
    right: 15%;
    height: 1px;
    background: linear-gradient(90deg,
        transparent,
        rgba(34, 197, 94, 0.25),
        transparent);
}

.output-panel::after {
    content: '';
    position: absolute;
    top: 0;
    left: 0;
    right: 0;
    bottom: 0;
    background: linear-gradient(180deg, rgba(34, 197, 94, 0.01) 0%, transparent 30%);
    pointer-events: none;
}

/* ═══ Footer ═══ */
.footer-block {
    text-align: center !important;
    margin-top: 36px !important;
    padding-top: 24px !important;
    border-top: 1px solid rgba(255, 255, 255, 0.03) !important;
}

.footer-block p {
    font-size: 11.5px !important;
    color: rgba(255, 255, 255, 0.18) !important;
    font-weight: 400 !important;
    letter-spacing: 0.02em !important;
}

.tech-row {
    display: flex !important;
    justify-content: center !important;
    gap: 10px !important;
    margin-top: 14px !important;
    flex-wrap: wrap !important;
}

.tech-tag {
    font-size: 10.5px !important;
    color: rgba(255, 255, 255, 0.22) !important;
    background: rgba(255, 255, 255, 0.025) !important;
    border: 1px solid rgba(255, 255, 255, 0.03) !important;
    padding: 5px 14px !important;
    border-radius: 8px !important;
    font-weight: 500 !important;
    font-family: 'JetBrains Mono', monospace !important;
    letter-spacing: 0.03em !important;
    transition: all 0.3s ease !important;
}

.tech-tag:hover {
    border-color: rgba(255, 255, 255, 0.08) !important;
    color: rgba(255, 255, 255, 0.35) !important;
}

/* ═══ macOS-style Scrollbar ═══ */
::-webkit-scrollbar {
    width: 5px !important;
}
::-webkit-scrollbar-track {
    background: transparent !important;
}
::-webkit-scrollbar-thumb {
    background: rgba(255, 255, 255, 0.1) !important;
    border-radius: 10px !important;
}
::-webkit-scrollbar-thumb:hover {
    background: rgba(255, 255, 255, 0.18) !important;
}

/* ═══ Remove Gradio Footer ═══ */
footer {
    display: none !important;
}

/* ═══ Entry Animation ═══ */
@keyframes slideUp {
    from { opacity: 0; transform: translateY(20px); }
    to { opacity: 1; transform: translateY(0); }
}

.main-container {
    animation: slideUp 0.7s cubic-bezier(0.16, 1, 0.3, 1) !important;
}

/* ═══ Row Gaps ═══ */
.gr-row {
    gap: 24px !important;
}

/* ═══ Selection Color ═══ */
::selection {
    background: rgba(139, 92, 246, 0.3) !important;
    color: #fff !important;
}
"""

# ──────────────────────────────────────────────
# Build the Aurora Noir Interface
# ──────────────────────────────────────────────
with gr.Blocks(title="AI Rewriter") as demo:

    # Main Container
    with gr.Column(elem_classes="main-container"):

        # ── Top Navigation Bar ──
        gr.HTML("""
            <div class="top-bar">
                <div class="logo-group">
                    <div class="logo-icon">✦</div>
                    <span class="logo-text">AI Rewriter</span>
                </div>
                <div class="status-badge">
                    <span class="status-dot"></span>
                    Model Active
                </div>
            </div>
        """)

        # ── Content Body ──
        with gr.Column(elem_classes="content-body"):

            # ── Hero Section ──
            gr.HTML("""
                <div class="hero-block">
                    <h1>Rewrite with <span class="gradient-word">Intelligence</span></h1>
                    <p class="tagline">
                        Transform any draft into polished, professional text.
                        Powered by a fine-tuned 1.5B parameter model running entirely on-device.
                    </p>
                    <div class="feature-chips">
                        <span class="feature-chip">🔒 100% Private</span>
                        <span class="feature-chip">⚡ On-Device</span>
                        <span class="feature-chip">🎯 Fine-Tuned</span>
                        <span class="feature-chip">∞ Unlimited</span>
                    </div>
                    <div class="divider-line"></div>
                </div>
            """)

            with gr.Row():

                # ── Left: Input Panel ──
                with gr.Column(scale=1):
                    gr.HTML("""
                        <div class="panel-label">
                            <div class="panel-label-icon input-icon">⬤</div>
                            <span class="panel-label-text">Input</span>
                        </div>
                    """)

                    with gr.Column(elem_classes="glass-card"):
                        instruction_input = gr.Textbox(
                            label="Instructions",
                            placeholder="e.g., Rewrite as a formal business email...",
                            lines=2,
                            show_label=True
                        )

                    with gr.Column(elem_classes="glass-card"):
                        text_input = gr.Textbox(
                            label="Original Text",
                            placeholder="Paste your draft here...",
                            lines=6,
                            show_label=True
                        )

                    submit_btn = gr.Button(
                        "✦  Generate Rewrite",
                        variant="primary"
                    )

                # ── Right: Output Panel ──
                with gr.Column(scale=1):
                    gr.HTML("""
                        <div class="panel-label">
                            <div class="panel-label-icon output-icon">⬤</div>
                            <span class="panel-label-text">Output</span>
                        </div>
                    """)

                    with gr.Column(elem_classes="output-panel"):
                        output_display = gr.Textbox(
                            label="Rewritten Text",
                            interactive=False,
                            lines=12,
                            show_label=True
                        )

            # ── Footer ──
            gr.HTML("""
                <div class="footer-block">
                    <p>Engineered for precision & privacy</p>
                    <div class="tech-row">
                        <span class="tech-tag">Qwen 1.5B</span>
                        <span class="tech-tag">Q4_K_M</span>
                        <span class="tech-tag">llama.cpp</span>
                        <span class="tech-tag">Free CPU</span>
                    </div>
                </div>
            """)

    # ── Connect Logic ──
    submit_btn.click(
        fn=rewrite_text,
        inputs=[instruction_input, text_input],
        outputs=output_display
    )

# Launch with CSS
demo.launch(css=aurora_css,ssr_mode=False)