app.py

"""Gradio ZeroGPU demo for Talkie — a 13B vintage language model from 1930.

Loads the instruction-tuned 1930 model for chat.
"""

from __future__ import annotations

import os
import sys

# Make the cloned talkie repo importable without installing.
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "talkie", "src"))

import gradio as gr
import spaces
import torch
from gradio.themes.base import Base
from gradio.themes.utils import colors, fonts, sizes

from talkie import Message, Talkie

# ---------------------------------------------------------------------------
# Model loading (root-module level, as required by ZeroGPU).
# ---------------------------------------------------------------------------

device = "cuda" if torch.cuda.is_available() else "cpu"

print("Loading talkie-1930-13b-it ...")
it_model = Talkie("talkie-1930-13b-it", device=device)
print("Model loaded.")


# ---------------------------------------------------------------------------
# Inference functions.
# ---------------------------------------------------------------------------


def _history_to_messages(
    history: list[dict], system_prompt: str | None
) -> list[Message]:
    msgs: list[Message] = []
    if system_prompt and system_prompt.strip():
        msgs.append(Message(role="system", content=system_prompt.strip()))
    for entry in history:
        role = entry["role"]
        if role in ("user", "assistant"):
            msgs.append(Message(role=role, content=entry["content"]))
    return msgs


@spaces.GPU(duration=120)
def chat_fn(
    message: str,
    history: list[dict],
    system_prompt: str,
    temperature: float,
    max_tokens: int,
    top_p: float,
    top_k: int,
):
    history = list(history) + [{"role": "user", "content": message}]
    messages = _history_to_messages(history, system_prompt)

    response = ""
    for token in it_model.chat_stream(
        messages,
        temperature=float(temperature),
        max_tokens=int(max_tokens),
        top_p=float(top_p) if top_p < 1.0 else None,
        top_k=int(top_k) if top_k > 0 else None,
    ):
        response += token
        yield response


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------

DESCRIPTION = """
<div class="masthead">
  <div class="masthead-rule"></div>
  <h1 class="masthead-title">Talkie</h1>
  <div class="masthead-rule"></div>
  <div class="masthead-meta">
    <span>VOL. XIII · No. up to 1930s</span>
    <span>A 13-BILLION-PARAMETER VINTAGE LANGUAGE MODEL</span>
    <span>TRAINED ON 260B TOKENS</span>
  </div>
  <div class="masthead-rule double"></div>
</div>

<p class="lede">
<span class="dropcap">T</span>HE <a href="https://github.com/talkie-lm/talkie">talkie</a>
team presents the <strong>Talkie 13B</strong> family of language
models — gentlemen and ladies of letters, trained exclusively upon
<em>pre-1931 English-language text</em>. Herein you may converse with
<code>talkie-1930-13b-it</code>, an instruction-tuned variant schooled in
pre-1931 etiquette manuals, letter-writing primers, encyclopedias, and poetry
collections, then refined by online DPO.</p>

<p class="byline">By Alec Radford, Nick Levine &amp; David Duvenaud —
<a href="https://talkie-lm.com/">read the dispatch</a>.</p>
""".strip()

CHAT_EXAMPLES = [
    "Write an essay predicting what life will be like in the year 1960.",
    "What were the causes of the French Revolution?",
    "Compose a short letter declining a dinner invitation politely.",
    "Explain the workings of the steam engine to a young pupil.",
    "What is the proper etiquette for greeting a lady on the street?",
]


# ---------------------------------------------------------------------------
# 1930s newspaper theme.
# ---------------------------------------------------------------------------


class Newsprint(Base):
    def __init__(self):
        super().__init__(
            primary_hue=colors.stone,
            secondary_hue=colors.amber,
            neutral_hue=colors.stone,
            spacing_size=sizes.spacing_md,
            radius_size=sizes.radius_none,
            text_size=sizes.text_md,
            font=(
                fonts.GoogleFont("IM Fell English"),
                "Georgia",
                "ui-serif",
                "serif",
            ),
            font_mono=(
                fonts.GoogleFont("Special Elite"),
                "Courier New",
                "ui-monospace",
                "monospace",
            ),
        )
        super().set(
            # Page = newsprint paper
            body_background_fill="#efe6d3",
            body_background_fill_dark="#1a1612",
            body_text_color="#1a1410",
            body_text_color_dark="#e8dec6",
            body_text_color_subdued="#5a4a3a",
            background_fill_primary="#f6efdc",
            background_fill_primary_dark="#221c16",
            background_fill_secondary="#ece2c8",
            background_fill_secondary_dark="#2b231b",
            # Hairline rules everywhere
            border_color_primary="#1a1410",
            border_color_primary_dark="#9a8a70",
            border_color_accent="#1a1410",
            block_border_width="1px",
            block_border_color="#1a1410",
            block_label_border_width="0px",
            input_border_width="1px",
            input_border_color="#1a1410",
            panel_border_width="1px",
            # Buttons — black ink on cream
            button_primary_background_fill="#1a1410",
            button_primary_background_fill_hover="#3a2c1a",
            button_primary_text_color="#f6efdc",
            button_primary_border_color="#1a1410",
            button_secondary_background_fill="#ece2c8",
            button_secondary_background_fill_hover="#dccfaf",
            button_secondary_text_color="#1a1410",
            button_secondary_border_color="#1a1410",
            # Block titles look like article headings
            block_title_text_color="#1a1410",
            block_title_text_color_dark="#e8dec6",
            block_title_text_weight="700",
            block_label_text_color="#5a4a3a",
            block_label_background_fill="transparent",
            # Shadow off — period correct
            shadow_drop="none",
            shadow_drop_lg="none",
            block_shadow="none",
            button_primary_shadow="none",
            # Inputs
            input_background_fill="#fbf6e6",
            input_background_fill_dark="#2b231b",
            input_placeholder_color="#8a7a64",
            # Slider in sepia
            slider_color="#3a2c1a",
            slider_color_dark="#c9b78f",
        )


newsprint = Newsprint()


CUSTOM_CSS = """
@import url('https://fonts.googleapis.com/css2?family=UnifrakturCook:wght@700&family=Playfair+Display:ital,wght@0,400;0,700;0,900;1,400&family=IM+Fell+English:ital@0;1&family=Special+Elite&display=swap');

/* ---------- Theme tokens (light/dark via CSS vars) ---------- */
.gradio-container {
    --paper: #efe6d3;
    --paper-light: #f6efdc;
    --paper-mid: #ece2c8;
    --paper-input: #fbf6e6;
    --ink: #1a1410;
    --ink-soft: #3a2c1a;
    --ink-faint: #5a4a3a;
    --ornament: #5a4a3a;
}
.dark .gradio-container,
.gradio-container.dark,
.dark.gradio-container {
    --paper: #1a1612;
    --paper-light: #221c16;
    --paper-mid: #2b231b;
    --paper-input: #2b231b;
    --ink: #e8dec6;
    --ink-soft: #d8c9a3;
    --ink-faint: #b8a988;
    --ornament: #c9b78f;
}

.gradio-container {
    max-width: 1100px !important;
    margin: 0 auto !important;
    background: var(--paper) !important;
    background-attachment: fixed !important;
    padding: 28px 36px !important;
    font-family: "IM Fell English", Georgia, serif;
    color: var(--ink) !important;
}

/* ---------- Masthead ---------- */
.masthead {
    text-align: center;
    margin: 4px 0 18px 0;
    color: var(--ink);
}
.masthead-title {
    font-family: "UnifrakturCook", "Playfair Display", Georgia, serif !important;
    font-weight: 700 !important;
    font-size: clamp(48px, 7vw, 84px) !important;
    line-height: 1 !important;
    letter-spacing: 1px;
    margin: 4px 0 4px 0 !important;
    color: var(--ink) !important;
}
.masthead-rule {
    border-top: 2px solid var(--ink);
    margin: 8px 0;
}
.masthead-rule.double {
    border-top: 1px solid var(--ink);
    border-bottom: 1px solid var(--ink);
    height: 4px;
    margin-top: 10px;
}
.masthead-meta {
    display: flex;
    justify-content: space-between;
    flex-wrap: wrap;
    gap: 8px 24px;
    font-family: "Playfair Display", Georgia, serif;
    font-variant: small-caps;
    letter-spacing: 1.5px;
    font-size: 12px;
    text-transform: uppercase;
    padding: 4px 0;
    color: var(--ink-soft);
}

/* ---------- Lede paragraph & dropcap ---------- */
.lede {
    font-family: "IM Fell English", Georgia, serif;
    font-size: 17px;
    line-height: 1.55;
    text-align: justify;
    hyphens: auto;
    margin: 14px 0 8px 0;
    color: var(--ink);
}
.dropcap {
    float: left;
    font-family: "UnifrakturCook", "Playfair Display", Georgia, serif;
    font-size: 58px;
    line-height: 0.9;
    padding: 6px 10px 0 0;
    color: var(--ink);
}
.byline {
    font-family: "Playfair Display", Georgia, serif;
    font-style: italic;
    text-align: center;
    font-size: 14px;
    margin: 4px 0 18px 0;
    color: var(--ink-soft);
    border-bottom: 3px double var(--ink);
    padding-bottom: 12px;
}
.lede a, .byline a {
    color: var(--ink);
    text-decoration: underline;
}

/* ---------- Section headings (block labels) ---------- */
.gradio-container h1, .gradio-container h2, .gradio-container h3 {
    font-family: "Playfair Display", Georgia, serif !important;
    font-weight: 900 !important;
    text-align: center;
    text-transform: uppercase;
    letter-spacing: 2px;
    color: var(--ink) !important;
}

label, .label-wrap span, span[data-testid="block-label"] {
    font-family: "Playfair Display", Georgia, serif !important;
    font-variant: small-caps !important;
    letter-spacing: 1.2px !important;
    color: var(--ink) !important;
}

/* ---------- Chatbot styled like newsprint columns ---------- */
.gradio-container .chatbot, .gradio-container [data-testid="chatbot"] {
    background: var(--paper-input) !important;
    border: 1px solid var(--ink) !important;
    border-top: 4px double var(--ink) !important;
    border-bottom: 4px double var(--ink) !important;
    box-shadow: none !important;
    color: var(--ink) !important;
}

/* Hide empty avatar slots that leave ghost boxes on the left */
.gradio-container .chatbot .avatar-container,
.gradio-container .chatbot [class*="avatar"] {
    display: none !important;
}

/* Message text — newspaper serif */
.gradio-container .chatbot .message,
.gradio-container .chatbot .message *,
.gradio-container .chatbot [data-testid="user"],
.gradio-container .chatbot [data-testid="user"] *,
.gradio-container .chatbot [data-testid="bot"],
.gradio-container .chatbot [data-testid="bot"] * {
    font-family: "IM Fell English", Georgia, serif !important;
    font-size: 17px !important;
    line-height: 1.6 !important;
    color: var(--ink) !important;
}

/* User bubble = ledger card */
.gradio-container .chatbot [data-testid="user"],
.gradio-container .chatbot .message.user {
    background: var(--paper-mid) !important;
    border: 1px solid var(--ink) !important;
    border-radius: 0 !important;
    box-shadow: 2px 2px 0 var(--ink) !important;
    font-style: italic !important;
}

/* Bot bubble = printed paragraph card */
.gradio-container .chatbot [data-testid="bot"],
.gradio-container .chatbot .message.bot {
    background: var(--paper-input) !important;
    border: 1px solid var(--ink) !important;
    border-radius: 0 !important;
    box-shadow: 2px 2px 0 var(--ink) !important;
}

/* ---------- Buttons ---------- */
button.primary, button[variant="primary"] {
    font-family: "Playfair Display", Georgia, serif !important;
    text-transform: uppercase !important;
    letter-spacing: 2px !important;
    font-weight: 700 !important;
    border-radius: 0 !important;
    border: 1px solid var(--ink) !important;
    background: var(--ink) !important;
    color: var(--paper-light) !important;
    box-shadow: 2px 2px 0 var(--ink) !important;
}
button.primary:hover, button[variant="primary"]:hover {
    transform: translate(1px, 1px);
    box-shadow: 1px 1px 0 var(--ink) !important;
}
button.secondary, button[variant="secondary"] {
    font-family: "Playfair Display", Georgia, serif !important;
    text-transform: uppercase !important;
    letter-spacing: 1.5px !important;
    border-radius: 0 !important;
    border: 1px solid var(--ink) !important;
    background: var(--paper-mid) !important;
    color: var(--ink) !important;
}

/* ---------- Inputs ---------- */
textarea, input[type=text], input[type=number] {
    font-family: "IM Fell English", Georgia, serif !important;
    background: var(--paper-input) !important;
    color: var(--ink) !important;
    border-radius: 0 !important;
}

/* ---------- Examples panel ---------- */
.gradio-container .examples-table, .gradio-container [data-testid="examples"] {
    background: var(--paper-mid) !important;
    border: 1px solid var(--ink) !important;
    border-radius: 0 !important;
}
.gradio-container .examples-table button {
    font-family: "IM Fell English", Georgia, serif !important;
    font-style: italic;
    background: var(--paper-input) !important;
    border: 1px solid var(--ink) !important;
    border-radius: 0 !important;
    color: var(--ink) !important;
}

/* ---------- Sliders ---------- */
input[type=range] {
    accent-color: var(--ink-soft);
}

/* ---------- Tabs / panels ---------- */
.gradio-container .panel, .gradio-container .form, .gradio-container .block {
    border-radius: 0 !important;
}

/* ---------- Footer hide ---------- */
footer { display: none !important; }

/* ---------- Mobile: remove excess padding ---------- */
@media (max-width: 768px) {
    .gradio-container {
        padding: 12px 8px !important;
    }
    .masthead-meta {
        justify-content: center;
        font-size: 10px;
        gap: 4px 12px;
    }
    .lede {
        font-size: 15px;
        text-align: left;
    }
}
"""


with gr.Blocks(title="Talkie 1930", theme=newsprint, css=CUSTOM_CSS) as demo:
    gr.HTML(DESCRIPTION)

    chatbot = gr.Chatbot(
        height=350,
        label="talkie-1930-13b-it",
        avatar_images=(None, None),
    )
    chat_input = gr.Textbox(
        placeholder="Ask the 1930 model anything…",
        show_label=False,
        submit_btn=True,
    )

    with gr.Accordion("Advanced settings", open=False):
        chat_system = gr.Textbox(
            label="System prompt (optional)",
            value="",
            lines=3,
            placeholder="e.g. You are a learned gentleman of the 1920s.",
        )
        with gr.Row():
            chat_temp = gr.Slider(
                0.0, 2.0, value=0.7, step=0.05, label="Temperature"
            )
            chat_max = gr.Slider(
                16, 1024, value=400, step=16, label="Max new tokens"
            )
        with gr.Row():
            chat_topp = gr.Slider(
                0.1, 1.0, value=1.0, step=0.05, label="Top-p (1.0 = off)"
            )
            chat_topk = gr.Slider(
                0, 200, value=0, step=1, label="Top-k (0 = off)"
            )

    gr.Examples(
        examples=[[ex] for ex in CHAT_EXAMPLES],
        inputs=[chat_input],
        label="Example prompts",
    )

    def _user_submit(user_msg, history):
        if not user_msg:
            return "", history
        history = history + [{"role": "user", "content": user_msg}]
        return "", history

    def _bot_reply(history, system_prompt, temperature, max_tokens, top_p, top_k):
        if not history or history[-1]["role"] != "user":
            yield history
            return
        user_msg = history[-1]["content"]
        prior = history[:-1]
        history = history + [{"role": "assistant", "content": ""}]
        for partial in chat_fn(
            user_msg, prior, system_prompt, temperature, max_tokens, top_p, top_k
        ):
            history[-1] = {"role": "assistant", "content": partial}
            yield history

    chat_state_inputs = [chat_system, chat_temp, chat_max, chat_topp, chat_topk]
    chat_input.submit(
        _user_submit,
        [chat_input, chatbot],
        [chat_input, chatbot],
        queue=False,
    ).then(
        _bot_reply,
        [chatbot, *chat_state_inputs],
        chatbot,
    )

    with gr.Row():
        clear_btn = gr.Button("Clear conversation", variant="secondary")
        clear_btn.click(lambda: [], None, chatbot, queue=False)


if __name__ == "__main__":
    demo.launch()