# app.py ─ Hermes Chatbot · HuggingFace Space
# Model : Havoc999/tiny-openhermes (TinyLlama 1.1B fine-tuned on OpenHermes-2.5)
# Run   : python app.py   OR deploy as a HuggingFace Space (set hardware → CPU Basic or T4)

import torch
import gradio as gr
from threading import Thread
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer

# ── Config ────────────────────────────────────────────────────────────────────
MODEL_ID = "Havoc999/tiny-openhermes"
DEVICE   = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE    = torch.float16 if torch.cuda.is_available() else torch.float32

# Generation params tuned for a 1.1B model:
#   MMLU 23% / HellaSwag 59% / ARC 30% → cap tokens, punish repetition hard,
#   keep temperature modest so it doesn't hallucinate wildly.
GEN = dict(
    max_new_tokens       = 256,   # short focused answers > wandering essays
    temperature          = 0.72,  # low enough to reduce confabulation
    top_p                = 0.90,  # nucleus keeps diversity without wild tokens
    top_k                = 45,    # hard vocab cap stops the long tail
    repetition_penalty   = 1.20,  # tiny models loop — this is the main guard
    no_repeat_ngram_size = 3,     # second line of defence against phrase loops
    do_sample            = True,
)

# ── Identity seed injected as a synthetic first exchange ─────────────────────
# The model was trained without <|system|> tokens, so we plant its identity
# via a Q&A pair that silently prepends every conversation.
SEED = (
    "Who are you?",
    "I'm Hermes — a compact AI assistant built on TinyLlama 1.1B and fine-tuned "
    "on OpenHermes-2.5. I can answer questions, help with writing and coding, explain "
    "concepts, and hold a conversation. I'm a small model, so clear focused questions "
    "get the best results.",
)

# ── Load model ────────────────────────────────────────────────────────────────
print(f"▸ Loading {MODEL_ID}  on {DEVICE} ({DTYPE}) …")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=True)
if tokenizer.pad_token is None:
    tokenizer.pad_token    = tokenizer.eos_token
    tokenizer.pad_token_id = tokenizer.eos_token_id

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype       = DTYPE,
    device_map        = "auto",
    low_cpu_mem_usage = True,
)
model.eval()
print("▸ Model ready ✓")

# ── Prompt builder ────────────────────────────────────────────────────────────
def build_prompt(history: list, user_msg: str) -> str:
    """Convert [(user,bot),…] history + new user message → TinyLlama chat format."""
    text = ""
    # prepend identity seed so model always knows its name & purpose
    for u, b in [(SEED[0], SEED[1])] + history:
        text += f"<|user|>\n{u}</s>\n<|assistant|>\n{b}</s>\n"
    text += f"<|user|>\n{user_msg}</s>\n<|assistant|>\n"
    return text

def clean(text: str) -> str:
    """Strip any leaked role tags and trailing whitespace from model output."""
    for tag in ("<|user|>", "<|assistant|>", "</s>", "<s>"):
        text = text.split(tag)[0]
    return text.strip()

# ── Step 1: user click / enter → clear input, append pending row ──────────────
def user_step(msg: str, history: list):
    """Immediately blank the textbox and add the user turn (bot slot = None)."""
    if not msg or not msg.strip():
        return "", history
    return "", history + [[msg.strip(), None]]

# ── Step 2: stream the model's response into the pending row ──────────────────
def bot_step(history: list):
    """Stream tokens into history[-1][1]; yield after each token for live display."""
    if not history or history[-1][1] is not None:
        yield history
        return

    user_msg       = history[-1][0]
    history[-1][1] = ""
    prior          = history[:-1]

    prompt = build_prompt(prior, user_msg)
    inputs = tokenizer(
        prompt,
        return_tensors = "pt",
        truncation     = True,
        max_length     = 2048 - GEN["max_new_tokens"],
    ).to(DEVICE)

    streamer = TextIteratorStreamer(
        tokenizer,
        skip_prompt        = True,
        skip_special_tokens= True,
        timeout            = 60.0,
    )
    Thread(
        target = model.generate,
        kwargs = dict(
            **inputs,
            streamer     = streamer,
            pad_token_id = tokenizer.pad_token_id,
            eos_token_id = tokenizer.eos_token_id,
            **GEN,
        ),
        daemon = True,
    ).start()

    for token in streamer:
        history[-1][1] += token
        # hard-stop if a role tag leaks into the generation
        if any(t in history[-1][1] for t in ("<|user|>", "<|assistant|>")):
            history[-1][1] = clean(history[-1][1])
            yield history
            return
        yield history

    history[-1][1] = clean(history[-1][1])
    yield history


# ═══════════════════════════════════════════════════════════════════════════════
#  CSS — Hermes / Mercury aesthetic
#  Palette: deep navy-black · cold steel-blue · quicksilver text · gold accent
# ═══════════════════════════════════════════════════════════════════════════════
CSS = """
/* ── Fonts ── */
@import url('https://fonts.googleapis.com/css2?family=DM+Sans:wght@300;400;500;600;700&display=swap');

/* ── Tokens ── */
:root {
  --abyss     : #07090f;
  --surface   : #0c0f1a;
  --surface-2 : #111827;
  --border    : #1a2540;
  --border-hi : #253554;
  --text      : #bfcbe8;
  --text-muted: #3d5070;
  --text-dim  : #5b7099;
  --accent    : #3b82f6;
  --accent-hi : #60a5fa;
  --gold      : #c9a227;
  --gold-dim  : #7a6218;
  --user-a    : #1e3a5f;
  --user-b    : #1a4a7a;
  --bot-bg    : #0c1220;
  --danger    : #ef4444;
}

/* ── Base ── */
body,
.gradio-container,
.gradio-container > .main {
  background : var(--abyss) !important;
  font-family: 'DM Sans', system-ui, -apple-system, sans-serif !important;
  color      : var(--text) !important;
}

/* hide Gradio's built-in footer */
footer { display: none !important; }

/* ── Outer wrapper ── */
.gradio-container {
  max-width  : 800px !important;
  margin     : 0 auto !important;
  padding    : 1rem !important;
}

/* ── Header ── */
#hdr {
  text-align : center;
  padding    : 2rem 1.5rem 1.6rem;
  background : linear-gradient(160deg, #0c1220 0%, #0e1630 55%, #0a1428 100%);
  border     : 1px solid var(--border);
  border-radius: 20px;
  margin-bottom: 1rem;
  position   : relative;
  overflow   : hidden;
}

/* aurora shimmer behind the header — the signature element */
#hdr::before {
  content   : '';
  position  : absolute;
  inset     : 0;
  background: radial-gradient(ellipse 70% 60% at 50% -10%,
                rgba(59,130,246,0.12) 0%,
                rgba(201,162,39,0.04) 60%,
                transparent 100%);
  pointer-events: none;
}

#hdr .eyebrow {
  font-size    : 0.68rem;
  letter-spacing: 0.22em;
  text-transform: uppercase;
  color        : var(--gold);
  margin-bottom: 0.6rem;
}

#hdr h1 {
  font-size  : 2.4rem;
  font-weight: 700;
  letter-spacing: -0.02em;
  color      : var(--text);
  margin     : 0 0 0.2rem;
  line-height: 1.1;
}

/* the single gold accent: a centred divider ◆ */
#hdr .divider {
  display    : flex;
  align-items: center;
  gap        : 0.6rem;
  justify-content: center;
  margin     : 0.75rem 0 0.6rem;
}
#hdr .divider::before,
#hdr .divider::after {
  content   : '';
  flex      : 1;
  max-width : 80px;
  height    : 1px;
  background: linear-gradient(90deg, transparent, var(--gold-dim));
}
#hdr .divider::after {
  background: linear-gradient(90deg, var(--gold-dim), transparent);
}
#hdr .divider span {
  color    : var(--gold);
  font-size: 0.65rem;
}

#hdr p {
  color    : var(--text-dim);
  font-size: 0.8rem;
  margin   : 0;
  letter-spacing: 0.04em;
}

/* ── Chatbot container ── */
#chatbox {
  background   : var(--surface) !important;
  border       : 1px solid var(--border) !important;
  border-radius: 16px !important;
  height       : 520px !important;
  overflow-y   : auto !important;
}

/* Gradio 4.x: target both selector patterns for max compatibility */

/* User bubbles */
#chatbox .message-row.user-row .bubble-wrap,
#chatbox .bubble-wrap.user,
#chatbox div[data-testid="user"] {
  background   : linear-gradient(135deg, var(--user-a), var(--user-b)) !important;
  border-radius: 18px 18px 4px 18px !important;
  border       : 1px solid #253c60 !important;
  margin-left  : auto !important;
  max-width    : 78% !important;
  box-shadow   : 0 2px 10px rgba(30,74,122,0.35) !important;
}

/* Bot bubbles */
#chatbox .message-row.bot-row .bubble-wrap,
#chatbox .bubble-wrap.bot,
#chatbox div[data-testid="bot"] {
  background   : var(--bot-bg) !important;
  border-radius: 18px 18px 18px 4px !important;
  border       : 1px solid var(--border-hi) !important;
  margin-right : auto !important;
  max-width    : 78% !important;
}

/* Inner message text */
#chatbox .message,
#chatbox .bubble-wrap p,
#chatbox .bubble-wrap .prose {
  color    : var(--text) !important;
  font-size: 0.93rem !important;
  line-height: 1.6 !important;
  padding  : 0.65rem 0.9rem !important;
}

/* user text slightly brighter */
#chatbox .bubble-wrap.user .message,
#chatbox .message-row.user-row .bubble-wrap p,
#chatbox .message-row.user-row .bubble-wrap .prose {
  color: #dce8ff !important;
}

/* streaming cursor on the last bot message */
@keyframes blink { 0%,100%{opacity:1} 50%{opacity:0} }
#chatbox .bot .message:last-child::after,
#chatbox .bubble-wrap.bot:last-of-type p:last-child::after {
  content  : " ▍";
  color    : var(--accent-hi);
  animation: blink 0.9s infinite;
  font-size: 0.85em;
}

/* empty state */
#chatbox .empty {
  color    : var(--text-muted) !important;
  font-size: 0.85rem !important;
}

/* ── Input row ── */
#input-row {
  background   : var(--surface) !important;
  border       : 1px solid var(--border) !important;
  border-radius: 14px !important;
  padding      : 0.55rem 0.55rem !important;
  margin-top   : 0.75rem !important;
  align-items  : flex-end !important;
  gap          : 0.45rem !important;
}

/* textbox wrapper */
#msg-box {
  flex    : 1 !important;
  min-width: 0 !important;
}
#msg-box label { display: none !important; }   /* hide the "Textbox" label */
#msg-box textarea {
  background : #080c16 !important;
  color      : var(--text) !important;
  border     : 1px solid var(--border-hi) !important;
  border-radius: 9px !important;
  font-size  : 0.93rem !important;
  font-family: inherit !important;
  padding    : 0.6rem 0.85rem !important;
  resize     : none !important;
  min-height : 42px !important;
  line-height: 1.5 !important;
  outline    : none !important;
  transition : border-color 0.2s, box-shadow 0.2s !important;
}
#msg-box textarea:focus {
  border-color: var(--accent) !important;
  box-shadow  : 0 0 0 3px rgba(59,130,246,0.15) !important;
}
#msg-box textarea::placeholder { color: var(--text-muted) !important; }

/* Send button */
#send-btn {
  background   : linear-gradient(135deg, #1d4ed8, #2563eb) !important;
  border       : 1px solid #3b82f6 !important;
  border-radius: 9px !important;
  color        : #fff !important;
  font-weight  : 600 !important;
  font-size    : 0.88rem !important;
  letter-spacing: 0.03em !important;
  height       : 42px !important;
  min-width    : 78px !important;
  padding      : 0 1.1rem !important;
  cursor       : pointer !important;
  white-space  : nowrap !important;
  transition   : all 0.16s ease !important;
  box-shadow   : 0 2px 14px rgba(37,99,235,0.4) !important;
}
#send-btn:hover  { opacity: 0.88 !important; transform: translateY(-1px) !important; }
#send-btn:active { opacity: 0.70 !important; transform: translateY(0px) !important; }

/* ── Clear button ── */
#bottom-row {
  margin-top : 0.4rem !important;
  display    : flex !important;
  justify-content: flex-end !important;
}
#clear-btn {
  background   : transparent !important;
  color        : var(--text-muted) !important;
  border       : 1px solid var(--border) !important;
  border-radius: 8px !important;
  font-size    : 0.78rem !important;
  padding      : 0.25rem 0.85rem !important;
  cursor       : pointer !important;
  transition   : all 0.18s !important;
  letter-spacing: 0.02em !important;
}
#clear-btn:hover {
  border-color: var(--danger) !important;
  color       : var(--danger) !important;
  background  : rgba(239,68,68,0.06) !important;
}

/* ── Footer note ── */
#note {
  text-align : center;
  color      : var(--text-muted);
  font-size  : 0.71rem;
  margin-top : 0.6rem;
  letter-spacing: 0.04em;
}

/* ── Scrollbar ── */
* { scrollbar-width: thin; scrollbar-color: var(--border-hi) var(--abyss); }
::-webkit-scrollbar       { width: 5px; }
::-webkit-scrollbar-track { background: var(--abyss); }
::-webkit-scrollbar-thumb { background: var(--border-hi); border-radius: 6px; }

/* ── Mobile ── */
@media (max-width: 600px) {
  .gradio-container { padding: 0.4rem !important; }
  #hdr h1           { font-size: 1.8rem !important; }
  #chatbox          { height: 390px !important; }
  #send-btn         { min-width: 60px !important; font-size: 0.8rem !important; }
}
"""


# ── UI ────────────────────────────────────────────────────────────────────────
with gr.Blocks(css=CSS, title="Hermes · Tiny Chat", theme=gr.themes.Base()) as demo:

    # ── Header ──
    gr.HTML("""
    <div id="hdr">
      <div class="eyebrow">A I &nbsp; A S S I S T A N T</div>
      <h1>Hermes</h1>
      <div class="divider"><span>◆</span></div>
      <p>TinyLlama 1.1B &nbsp;·&nbsp; OpenHermes-2.5 &nbsp;·&nbsp; Ask me anything</p>
    </div>
    """)

    # ── Chat display ──
    chatbot = gr.Chatbot(
        elem_id          = "chatbox",
        label            = "",
        bubble_full_width= False,
        height           = 520,
        show_label       = False,
        avatar_images    = (None, None),
        placeholder      = "<div style='color:#3d5070;font-size:0.85rem;text-align:center;padding-top:3rem;'>Send a message to start chatting with Hermes</div>",
    )

    # ── Input row ──
    with gr.Row(elem_id="input-row"):
        msg = gr.Textbox(
            elem_id     = "msg-box",
            placeholder = "Type a message…  (Enter to send)",
            show_label  = False,
            scale       = 9,
            container   = False,
            lines       = 1,
            max_lines   = 6,
            autofocus   = True,
        )
        send = gr.Button("Send ↑", elem_id="send-btn", scale=1, min_width=78)

    # ── Controls row ──
    with gr.Row(elem_id="bottom-row"):
        clear = gr.Button("✕  Clear", elem_id="clear-btn", size="sm")

    # ── Footer ──
    gr.HTML("""
    <p id="note">
      Hermes &nbsp;·&nbsp; 1.1B parameters &nbsp;·&nbsp;
      May hallucinate &nbsp;·&nbsp; Best for focused, clear questions
    </p>
    """)

    # ── Event wiring ──────────────────────────────────────────────────────────
    # user_step: fires on Enter / Send click
    #   → returns ("", updated_history)  →  blanks the textbox instantly
    # bot_step:  fires immediately after, streams tokens into history[-1][1]

    msg.submit(
        fn      = user_step,
        inputs  = [msg, chatbot],
        outputs = [msg, chatbot],
        queue   = False,
    ).then(
        fn      = bot_step,
        inputs  = [chatbot],
        outputs = [chatbot],
    )

    send.click(
        fn      = user_step,
        inputs  = [msg, chatbot],
        outputs = [msg, chatbot],
        queue   = False,
    ).then(
        fn      = bot_step,
        inputs  = [chatbot],
        outputs = [chatbot],
    )

    clear.click(
        fn      = lambda: ([], ""),
        inputs  = [],
        outputs = [chatbot, msg],
        queue   = False,
    )


demo.queue(max_size=6, default_concurrency_limit=1)
demo.launch(show_api=False)