Update app.py

app.py

@@ -1,28 +1,28 @@
-#MODEL_ID = "ibm-granite/granite-3.1-2b-instruct"
+import time
+import gradio as gr
+from openai import OpenAI
+import os
+import subprocess
+from huggingface_hub import snapshot_download
+
+# Model configuration
 MODEL_ID = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
 QUANT = "Q5_K_M"
 
-import subprocess
+# Utility functions
 def run_command(command, cwd=None):
-    """运行系统命令"""
+    """Run a system command."""
     result = subprocess.run(command, shell=True, cwd=cwd, text=True, capture_output=True)
     if result.returncode != 0:
-        print(f"命令执行失败: {command}")
-        print(f"错误信息: {result.stderr}")
+        print(f"Command failed: {command}")
+        print(f"Error: {result.stderr}")
         exit(result.returncode)
     else:
-        print(f"命令执行成功: {command}")
+        print(f"Command succeeded: {command}")
         print(result.stdout)
 
-run_command('pip install llama-cpp-python')
-
-import gradio as gr
-import os
-from llama_cpp import Llama
-from huggingface_hub import snapshot_download
-
 def setup_llama_cpp():
-    """克隆并编译llama.cpp仓库"""
+    """Clone and compile llama.cpp repository."""
     if not os.path.exists('llama.cpp'):
         run_command('git clone https://github.com/ggml-org/llama.cpp.git')
         os.chdir('llama.cpp')
@@ -32,7 +32,7 @@ def setup_llama_cpp():
         os.chdir('..')
 
 def setup_model(model_id):
-    """下载并转换模型为GGUF格式，返回量化模型路径"""
+    """Download and convert model to GGUF format, return quantized model path."""
     local_dir = model_id.split('/')[-1]
     if not os.path.exists(local_dir):
         snapshot_download(repo_id=model_id, local_dir=local_dir)
@@ -44,49 +44,217 @@ def setup_model(model_id):
         run_command(f'./llama.cpp/build/bin/llama-quantize ./{gguf_path} {quantized_path} {QUANT}')
     return quantized_path
 
-def chat_with_model(message, history, system_prompt, temperature, max_tokens, top_k, top_p):
-    """调用Llama模型生成回复"""
-    messages = [{"role": "system", "content": system_prompt}]
-    for user_msg, assistant_msg in history:
-        messages.append({"role": "user", "content": user_msg})
-        messages.append({"role": "assistant", "content": assistant_msg})
-    messages.append({"role": "user", "content": message})
-    stream = llm.create_chat_completion(
-        messages=messages,
-        stream=True,
-        temperature=temperature,
-        top_k=top_k,
-        top_p=top_p,
-        max_tokens=max_tokens,
-        stop=["<|im_end|>"]
+def start_llama_server(model_path):
+    """Start llama-server in the background."""
+    cmd = f'./llama.cpp/build/bin/llama-server --host 0.0.0.0 --port 8080 --model {model_path} --ctx-size 32768'
+    process = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+    # Give the server a moment to start
+    time.sleep(5)
+    return process
+
+# GUI-specific utilities (unchanged from your original)
+def format_time(seconds_float):
+    total_seconds = int(round(seconds_float))
+    hours = total_seconds // 3600
+    remaining_seconds = total_seconds % 3600
+    minutes = remaining_seconds // 60
+    seconds = remaining_seconds % 60
+    if hours > 0:
+        return f"{hours}h {minutes}m {seconds}s"
+    elif minutes > 0:
+        return f"{minutes}m {seconds}s"
+    else:
+        return f"{seconds}s"
+
+DESCRIPTION = '''
+# Duplicate the space for free private inference.
+## DeepSeek-R1 Distill Qwen-1.5B Demo
+A reasoning model trained using RL (Reinforcement Learning) that demonstrates structured reasoning capabilities.
+'''
+
+CSS = """
+.spinner { animation: spin 1s linear infinite; display: inline-block; margin-right: 8px; }
+@keyframes spin { from { transform: rotate(0deg); } to { transform: rotate(360deg); } }
+.thinking-summary { cursor: pointer; padding: 8px; background: #f5f5f5; border-radius: 4px; margin: 4px 0; }
+.thought-content { padding: 10px; background: #f8f9fa; border-radius: 4px; margin: 5px 0; }
+.thinking-container { border-left: 3px solid #facc15; padding-left: 10px; margin: 8px 0; background: #210c29; }
+details:not([open]) .thinking-container { border-left-color: #290c15; }
+details { border: 1px solid #e0e0e0 !important; border-radius: 8px !important; padding: 12px !important; margin: 8px 0 !important; transition: border-color 0.2s; }
+"""
+
+client = OpenAI(base_url="http://localhost:8080/v1", api_key="no-key-required")
+
+def user(message, history):
+    return "", history + [[message, None]]
+
+class ParserState:
+    __slots__ = ['answer', 'thought', 'in_think', 'start_time', 'last_pos', 'total_think_time']
+    def __init__(self):
+        self.answer = ""
+        self.thought = ""
+        self.in_think = False
+        self.start_time = 0
+        self.last_pos = 0
+        self.total_think_time = 0.0
+
+def parse_response(text, state):
+    buffer = text[state.last_pos:]
+    state.last_pos = len(text)
+    while buffer:
+        if not state.in_think:
+            think_start = buffer.find('<think>')
+            if think_start != -1:
+                state.answer += buffer[:think_start]
+                state.in_think = True
+                state.start_time = time.perf_counter()
+                buffer = buffer[think_start + 7:]
+            else:
+                state.answer += buffer
+                break
+        else:
+            think_end = buffer.find('</think>')
+            if think_end != -1:
+                state.thought += buffer[:think_end]
+                duration = time.perf_counter() - state.start_time
+                state.total_think_time += duration
+                state.in_think = False
+                buffer = buffer[think_end + 8:]
+            else:
+                state.thought += buffer
+                break
+    elapsed = time.perf_counter() - state.start_time if state.in_think else 0
+    return state, elapsed
+
+def format_response(state, elapsed):
+    answer_part = state.answer.replace('<think>', '').replace('</think>', '')
+    collapsible = []
+    collapsed = "<details open>"
+    if state.thought or state.in_think:
+        if state.in_think:
+            total_elapsed = state.total_think_time + elapsed
+            formatted_time = format_time(total_elapsed)
+            status = f"🌀 Thinking for {formatted_time}"
+        else:
+            formatted_time = format_time(state.total_think_time)
+            status = f"✅ Thought for {formatted_time}"
+            collapsed = "<details>"
+        collapsible.append(
+            f"{collapsed}<summary>{status}</summary>\n\n<div class='thinking-container'>\n{state.thought}\n</div>\n</details>"
+        )
+    return collapsible, answer_part
+
+def generate_response(history, temperature, top_p, max_tokens, active_gen):
+    messages = [
+        {"role": "system", "content": "You are a helpful assistant."},
+        *[{"role": "user" if i % 2 == 0 else "assistant", "content": msg or ""} 
+          for i, (user_msg, assistant_msg) in enumerate(history[:-1])],
+        {"role": "user", "content": history[-1][0]}
+    ]
+    full_response = ""
+    state = ParserState()
+    try:
+        stream = client.chat.completions.create(
+            model="",  # Model name not needed with llama-server
+            messages=messages,
+            temperature=temperature,
+            top_p=top_p,
+            max_tokens=max_tokens,
+            stream=True
+        )
+        for chunk in stream:
+            if not active_gen[0]:
+                break
+            if chunk.choices[0].delta.content:
+                full_response += chunk.choices[0].delta.content
+                state, elapsed = parse_response(full_response, state)
+                collapsible, answer_part = format_response(state, elapsed)
+                history[-1][1] = "\n\n".join(collapsible + [answer_part])
+                yield history
+        state, elapsed = parse_response(full_response, state)
+        collapsible, answer_part = format_response(state, elapsed)
+        history[-1][1] = "\n\n".join(collapsible + [answer_part])
+        yield history
+    except Exception as e:
+        history[-1][1] = f"Error: {str(e)}"
+        yield history
+    finally:
+        active_gen[0] = False
+
+# GUI setup
+with gr.Blocks(css=CSS) as demo:
+    gr.Markdown(DESCRIPTION)
+    active_gen = gr.State([False])
+    
+    chatbot = gr.Chatbot(
+        elem_id="chatbot",
+        height=500,
+        show_label=False,
+        render_markdown=True
     )
-    response = ""
-    for chunk in stream:
-        if "choices" in chunk and chunk["choices"]:
-            text = chunk["choices"][0].get("delta", {}).get("content", "")
-            response += text
-            yield response
+
+    with gr.Row():
+        msg = gr.Textbox(
+            label="Message",
+            placeholder="Type your message...",
+            container=False,
+            scale=4
+        )
+        submit_btn = gr.Button("Send", variant='primary', scale=1)
+    
+    with gr.Column(scale=2):
+        with gr.Row():
+            clear_btn = gr.Button("Clear", variant='secondary')
+            stop_btn = gr.Button("Stop", variant='stop')
+        
+        with gr.Accordion("Parameters", open=False):
+            temperature = gr.Slider(minimum=0.1, maximum=1.5, value=0.6, label="Temperature")
+            top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, label="Top-p")
+            max_tokens = gr.Slider(minimum=2048, maximum=32768, value=4096, step=64, label="Max Tokens")
+
+    gr.Examples(
+        examples=[
+            ["How many r's are in the word strawberry?"],
+            ["Write 10 funny sentences that end in a fruit!"],
+            ["Let’s play word chains! I’ll start: PIZZA. Your turn! Next word must start with… A!"]
+        ],
+        inputs=msg,
+        label="Example Prompts"
+    )
+    
+    submit_event = submit_btn.click(
+        user, [msg, chatbot], [msg, chatbot], queue=False
+    ).then(
+        lambda: [True], outputs=active_gen
+    ).then(
+        generate_response, [chatbot, temperature, top_p, max_tokens, active_gen], chatbot
+    )
+    
+    msg.submit(
+        user, [msg, chatbot], [msg, chatbot], queue=False
+    ).then(
+        lambda: [True], outputs=active_gen
+    ).then(
+        generate_response, [chatbot, temperature, top_p, max_tokens, active_gen], chatbot
+    )
+    
+    stop_btn.click(
+        lambda: [False], None, active_gen, cancels=[submit_event]
+    )
+    
+    clear_btn.click(lambda: None, None, chatbot, queue=False)
 
 if __name__ == "__main__":
+    # Install dependencies
+    run_command('pip install llama-cpp-python openai')
     setup_llama_cpp()
     MODEL_PATH = setup_model(MODEL_ID)
-    llm = Llama(
-        model_path=MODEL_PATH,
-        verbose=False,
-        n_threads=4,
-        n_ctx=32768,
-        eos_token_id=151643,  # 设置正确的 EOS token ID，确保模型能正确结束
-    )
-    gr.ChatInterface(
-        fn=chat_with_model,
-        title="Llama GGUF Chatbot",
-        description="使用Llama GGUF量化模型进行推理",
-        additional_inputs_accordion=gr.Accordion(label="⚙️ 参数设置", open=False),
-        additional_inputs=[
-            gr.Textbox("You are a helpful assistant.", label="System Prompt"),
-            gr.Slider(0, 1, 0.6, label="Temperature"),
-            gr.Slider(100, 4096, 1000, label="Max Tokens"),
-            gr.Slider(1, 100, 40, label="Top K"),
-            gr.Slider(0, 1, 0.85, label="Top P"),
-        ],
-    ).queue().launch()
+    
+    # Start llama-server
+    server_process = start_llama_server(MODEL_PATH)
+    try:
+        # Launch GUI
+        demo.launch(server_name="0.0.0.0", server_port=7860)
+    finally:
+        # Cleanup: terminate the server process when the GUI is closed
+        server_process.terminate()
+        server_process.wait()