Update app.py

app.py

@@ -1,23 +1,53 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
+import time
 
 client = InferenceClient("lambdaindie/lambdai")
 
 css = """
-.thinking-html {
+body {
+    background-color: #000000;
+    color: #e0e0e0;
+    font-family: 'JetBrains Mono', monospace;
+}
+.gr-button {
+    background: linear-gradient(to right, #2a2a2a, #1f1f1f);
+    color: white;
+    border-radius: 10px;
+    padding: 8px 16px;
+    font-weight: bold;
+    font-family: 'JetBrains Mono', monospace;
+}
+.gr-button:hover {
+    background: #333;
+}
+.gr-textbox textarea {
+    background-color: #181818 !important;
+    color: #fff !important;
+    font-family: 'JetBrains Mono', monospace;
+    border-radius: 8px;
+}
+.gr-chat-message {
+    font-family: 'JetBrains Mono', monospace;
+}
+.markdown-think {
+    background-color: #000000;
+    border-left: 4px solid #555;
+    padding: 10px;
+    margin-bottom: 8px;
+    font-style: italic;
+    white-space: pre-wrap;
+    font-family: 'JetBrains Mono', monospace;
     display: flex;
     align-items: center;
-    gap: 8px;
-    color: #666;
-    font-style: italic;
-    margin-bottom: 5px;
-    animation: pulse 1.5s infinite;
+    gap: 10px;
+    animation: pulse 1.5s infinite ease-in-out;
 }
 .loader {
     width: 14px;
     height: 14px;
-    border: 2px solid #ccc;
-    border-top: 2px solid #666;
+    border: 2px solid #888;
+    border-top: 2px solid #e0e0e0;
     border-radius: 50%;
     animation: spin 1s linear infinite;
 }
@@ -25,41 +55,78 @@ css = """
     to { transform: rotate(360deg); }
 }
 @keyframes pulse {
-    0% { opacity: 1; transform: scale(1); }
-    50% { opacity: 0.6; transform: scale(1.05); }
-    100% { opacity: 1; transform: scale(1); }
+    0% { opacity: 0.6; }
+    50% { opacity: 1.0; }
+    100% { opacity: 0.6; }
 }
 """
 
-def respond(message, history):
-    thinking = (
-        "<div class='thinking-html'>"
-        "<div class='loader'></div>"
-        "Thinking..."
-        "</div>"
-    )
-    yield history + [[message, thinking]]
+def respond(message, history, system_message, max_tokens, temperature, top_p):
+    messages = [{"role": "system", "content": system_message}] if system_message else []
+
+    for user, assistant in history:
+        if user:
+            messages.append({"role": "user", "content": user})
+        if assistant:
+            messages.append({"role": "assistant", "content": assistant})
+
+    thinking_prompt = messages + [
+        {
+            "role": "user",
+            "content": f"{message}\n\nThink step-by-step before answering."
+        }
+    ]
+
+    reasoning = ""
+    yield '<div class="markdown-think"><div class="loader"></div>Thinking...</div>'
 
-    prompt = f"Think step by step and explain your reasoning before answering:\n\n{message}"
-    response = client.chat_completion([{"role": "user", "content": prompt}], stream=False)
-    output = response['choices'][0]['message']['content']
+    for chunk in client.chat_completion(
+        thinking_prompt,
+        max_tokens=max_tokens,
+        stream=True,
+        temperature=temperature,
+        top_p=top_p,
+    ):
+        token = chunk.choices[0].delta.content or ""
+        reasoning += token
+        styled_thought = f'<div class="markdown-think"><div class="loader"></div>{reasoning.strip()}</div>'
+        yield styled_thought
 
-    if "\n\n" in output:
-        reasoning, answer = output.split("\n\n", 1)
-    else:
-        reasoning, answer = "No reasoning provided.", output
+    time.sleep(0.5)
 
-    reasoning_md = f"> {reasoning.strip()}"
-    final = f"{reasoning_md}\n\n{answer.strip()}"
-    yield history + [[message, final]]
+    final_prompt = messages + [
+        {"role": "user", "content": message},
+        {"role": "assistant", "content": reasoning.strip()},
+        {"role": "user", "content": "Now answer based on your reasoning above."}
+    ]
 
-with gr.Blocks(css=css) as demo:
-    gr.Markdown("## Lambdai-v1-1B")
-    chatbot = gr.Chatbot()
-    msg = gr.Textbox(label="Message")
-    send = gr.Button("Send")
+    final_answer = ""
+    for chunk in client.chat_completion(
+        final_prompt,
+        max_tokens=max_tokens,
+        stream=True,
+        temperature=temperature,
+        top_p=top_p,
+    ):
+        token = chunk.choices[0].delta.content or ""
+        final_answer += token
+        yield final_answer.strip()
 
-    send.click(respond, [msg, chatbot], chatbot)
-    msg.submit(respond, [msg, chatbot], chatbot)
+demo = gr.ChatInterface(
+    fn=respond,
+    title="λambdAI",
+    theme=gr.themes.Base(primary_hue="gray"),
+    css=css,
+    additional_inputs=[
+        gr.Textbox(
+            value="You are a concise, logical AI that explains its reasoning clearly before answering.",
+            label="System Message"
+        ),
+        gr.Slider(64, 2048, value=512, step=1, label="Max Tokens"),
+        gr.Slider(0.1, 2.0, value=0.7, step=0.1, label="Temperature"),
+        gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-p")
+    ]
+)
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()