Refactor

app.py

@@ -20,40 +20,11 @@ DEFAULT_SYSTEM_PROMPT = (
     "ask for clarification only when needed, and do not invent details."
 )
 
-CSS = """
-.gradio-container img,
-.gradio-container video,
-.gradio-container canvas {
-    max-width: 100%;
-    max-height: 360px;
-    object-fit: contain;
-}
-.gradio-container button svg,
-.gradio-container [role="button"] svg,
-.gradio-container label svg,
-.gradio-container .icon svg {
-    width: 1.1rem !important;
-    height: 1.1rem !important;
-    max-width: 1.1rem !important;
-    max-height: 1.1rem !important;
-}
-footer {
-    display: none !important;
-}
-"""
 
 @lru_cache(maxsize=1)
 def get_llm() -> Llama:
-    print(
-        f"Loading GGUF model {MODEL_REPO}/{MODEL_FILE} "
-        f"with n_ctx={N_CTX}, n_batch={N_BATCH}, n_threads={N_THREADS}",
-        flush=True,
-    )
-    model_path = hf_hub_download(
-        repo_id=MODEL_REPO,
-        filename=MODEL_FILE,
-        local_dir=os.getenv("MODEL_DIR") or None,
-    )
+    print(f"Loading model {MODEL_REPO}/{MODEL_FILE}", flush=True)
+    model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)
     print(f"Model file ready: {model_path}", flush=True)
 
     return Llama(
@@ -73,8 +44,7 @@ def trim_history(history: list[Any]) -> list[Any]:
     if not history:
         return []
 
-    max_messages = MAX_HISTORY_TURNS * 2
-    return history[-max_messages:]
+    return history[-(MAX_HISTORY_TURNS * 2) :]
 
 
 def build_messages(
@@ -82,7 +52,7 @@ def build_messages(
     history: list[Any],
     system_message: str,
 ) -> list[dict[str, str]]:
-    messages = [{"role": "system", "content": system_message.strip()}]
+    messages = [{"role": "system", "content": system_message}]
 
     for item in trim_history(history):
         if isinstance(item, dict):
@@ -101,81 +71,53 @@ def build_messages(
     return messages
 
 
-def generate_reply(
-    message: str,
-    history: list[Any],
-    system_message: str,
-    max_tokens: int,
-    temperature: float,
-    top_p: float,
-    repeat_penalty: float,
-) -> Any:
+def respond(
+    message,
+    history: list[dict[str, str]],
+    system_message,
+    max_tokens,
+    temperature,
+    top_p,
+):
     if not message.strip():
         return
 
     print(f"Received chat request: {message[:120]!r}", flush=True)
-    yield "Request received. Loading the local model if needed..."
-
-    try:
-        llm = get_llm()
-        messages = build_messages(message, history, system_message or DEFAULT_SYSTEM_PROMPT)
-        output = ""
-        yield "Model ready. Generating..."
-
-        stream = llm.create_chat_completion(
-            messages=messages,
-            max_tokens=max_tokens,
-            temperature=temperature,
-            top_p=top_p,
-            repeat_penalty=repeat_penalty,
-            stream=True,
-        )
-
-        for chunk in stream:
-            delta = chunk["choices"][0].get("delta", {})
-            token = delta.get("content") or ""
-            if token:
-                output += token
-                yield output
-    except Exception as exc:
-        print(f"Generation failed: {exc!r}", flush=True)
-        yield f"Generation failed: {exc}"
+
+    llm = get_llm()
+    messages = build_messages(message, history, system_message or DEFAULT_SYSTEM_PROMPT)
+    response = ""
+
+    for chunk in llm.create_chat_completion(
+        messages=messages,
+        max_tokens=max_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        stream=True,
+    ):
+        token = chunk["choices"][0].get("delta", {}).get("content") or ""
+        if token:
+            response += token
+            yield response
 
 
 chatbot = gr.ChatInterface(
-    generate_reply,
+    respond,
     additional_inputs=[
-        gr.Textbox(
-            value=DEFAULT_SYSTEM_PROMPT,
-            label="System message",
-            lines=3,
-        ),
-        gr.Slider(
-            minimum=32,
-            maximum=1024,
-            value=384,
-            step=16,
-            label="Max new tokens",
-        ),
+        gr.Textbox(value=DEFAULT_SYSTEM_PROMPT, label="System message"),
+        gr.Slider(minimum=32, maximum=1024, value=384, step=16, label="Max new tokens"),
+        gr.Slider(minimum=0.0, maximum=1.5, value=0.7, step=0.05, label="Temperature"),
         gr.Slider(
-            minimum=0.0,
-            maximum=1.5,
-            value=0.7,
+            minimum=0.1,
+            maximum=1.0,
+            value=0.9,
             step=0.05,
-            label="Temperature",
-        ),
-        gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="Top-p"),
-        gr.Slider(
-            minimum=1.0,
-            maximum=1.3,
-            value=1.1,
-            step=0.01,
-            label="Repeat penalty",
+            label="Top-p (nucleus sampling)",
         ),
     ],
 )
 
-with gr.Blocks(css=CSS) as demo:
+with gr.Blocks() as demo:
     chatbot.render()