Update app.py

app.py

@@ -6,16 +6,14 @@ import time
 
 # --- Configuration ---
 MODEL_REPO_ID = "unsloth/DeepSeek-R1-0528-Qwen3-8B-GGUF"
-# IMPORTANT: Verify this filename exists in the "Files and versions" tab of the repo
-MODEL_FILENAME = "DeepSeek-R1-0528-Qwen3-8B-Q4_K_M.gguf"
-LOCAL_MODEL_PATH = f"./{MODEL_FILENAME}" # Download to current directory
+MODEL_FILENAME = "DeepSeek-R1-0528-Qwen3-8B-Q4_K_M.gguf" # IMPORTANT: Verify this filename
+LOCAL_MODEL_PATH = f"./{MODEL_FILENAME}"
 
-# LLM Llama Parameters (adjust based on your Space's resources)
-N_CTX = 2048          # Context window size. Default 2048. Max for this model is very large, but needs RAM.
-N_THREADS = None      # Number of CPU threads to use. None = Llama.cpp auto-detects.
-                      # On smaller CPU Spaces (e.g., 2-4 cores), explicitly setting N_THREADS=2 or N_THREADS=4 might be beneficial.
-N_GPU_LAYERS = 0      # Number of layers to offload to GPU. 0 for CPU-only. -1 for all possible.
-VERBOSE_LLAMA = True  # Enable verbose logging from llama.cpp
+# LLM Llama Parameters
+N_CTX = 2048
+N_THREADS = None
+N_GPU_LAYERS = 0
+VERBOSE_LLAMA = True
 
 # Generation parameters
 DEFAULT_MAX_NEW_TOKENS = 512
@@ -24,6 +22,10 @@ DEFAULT_TOP_P = 0.95
 DEFAULT_TOP_K = 40
 DEFAULT_REPEAT_PENALTY = 1.1
 
+# Qwen specific chat format elements (defined globally)
+IM_START_TOKEN = "<|im_start|>"
+IM_END_TOKEN = "<|im_end|>"
+
 # --- Global variable for the model ---
 llm = None
 
@@ -37,7 +39,7 @@ def download_model_if_needed():
                 repo_id=MODEL_REPO_ID,
                 filename=MODEL_FILENAME,
                 local_dir=".",
-                local_dir_use_symlinks=False, # Good practice for GGUF
+                local_dir_use_symlinks=False,
                 resume_download=True
             )
             end_time = time.time()
@@ -45,18 +47,17 @@ def download_model_if_needed():
             return True
         except Exception as e:
             print(f"Error downloading model: {e}")
-            print("Please ensure MODEL_FILENAME is correct and available in the repository.")
             print(f"Attempted to download: {MODEL_REPO_ID}/{MODEL_FILENAME}")
             return False
     else:
         print(f"Model file {MODEL_FILENAME} already exists.")
         return True
-    return False
+    return False # Should not be reached if logic is correct, but good for completeness
 
 # --- Model Loading ---
 def load_llm_model():
     global llm
-    if llm is None: # Load only if not already loaded
+    if llm is None:
         if not os.path.exists(LOCAL_MODEL_PATH):
             print("Model file not found. Cannot load.")
             return False
@@ -69,52 +70,56 @@ def load_llm_model():
                 n_threads=N_THREADS,
                 n_gpu_layers=N_GPU_LAYERS,
                 verbose=VERBOSE_LLAMA,
-                # logits_all=True, # Set to True if you need logits for all tokens (consumes more VRAM/RAM)
             )
             end_time = time.time()
             print(f"Model loaded successfully in {end_time - start_time:.2f} seconds.")
             return True
         except Exception as e:
             print(f"Error loading Llama model: {e}")
-            print("Ensure llama-cpp-python is installed correctly and the model file is valid.")
-            print(f"If you are on a resource-constrained environment (like free Hugging Face Spaces), "
-                  f"the model ({MODEL_FILENAME}, ~{os.path.getsize(LOCAL_MODEL_PATH)/(1024**3):.2f}GB) might be too large.")
-            print("Try reducing N_CTX or using a smaller model variant if available.")
-            llm = None # Ensure llm is None if loading failed
+            print(f"If on resource-constrained environment, model ({MODEL_FILENAME}, ~{os.path.getsize(LOCAL_MODEL_PATH)/(1024**3):.2f}GB if exists) might be too large.")
+            llm = None
             return False
     else:
         print("Model already loaded.")
         return True
 
-
 # --- Chat Function ---
 def predict(message, history, system_prompt, max_new_tokens, temperature, top_p, top_k, repeat_penalty):
     if llm is None:
         return "Model not loaded. Please check the logs."
 
-    # Qwen specific chat format elements
-    im_start_token = "<|im_start|>"
-    im_end_token = "<|im_end|>"
     # Common stop tokens for Qwen-like models
-    stop_tokens = [im_end_token, im_start_token + "user", im_start_token + "system", llm.token_eos()]
-
+    # Accessing global IM_START_TOKEN and IM_END_TOKEN
+    stop_tokens = [IM_END_TOKEN, IM_START_TOKEN + "user", IM_START_TOKEN + "system", llm.token_eos_str()] # Use string representation of EOS
 
-    # Format messages for llama_cpp
-    messages = []
+    messages_for_api = [] # Renamed to avoid conflict with Gradio's 'messages' type
     if system_prompt and system_prompt.strip():
-        messages.append({"role": "system", "content": system_prompt.strip()})
+        messages_for_api.append({"role": "system", "content": system_prompt.strip()})
+
+    # History for Gradio Chatbot with type="messages" is already in the correct format
+    # history will be a list of lists, where each inner list is [user_msg, ai_msg]
+    # or if type="messages", it's a list of dicts.
+    # Let's assume for now the input `history` from chatbot (when type="tuples")
+    # needs conversion if predict is called directly with such history.
+    # If chatbot type="messages", history is already List[Dict[str, str]]
+    # The `user_chat_fn` and `bot_response_fn` handle history in `messages` format for the chatbot.
+    # So, when `predict` is called by `bot_response_fn`, `history` is actually `history_for_predict`
+    # which is `chat_history[:-1]`. `chat_history` is a list of tuples.
+    # We need to convert this tuple-style history to OpenAI dict style for create_chat_completion.
+    
+    # The history passed from `bot_response_fn` (history_for_predict) is list of [user, assistant] tuples
+    for human_msg, ai_msg in history: # history here is history_for_predict from bot_response_fn
+        messages_for_api.append({"role": "user", "content": human_msg})
+        if ai_msg is not None:
+            messages_for_api.append({"role": "assistant", "content": ai_msg})
+    messages_for_api.append({"role": "user", "content": message})
 
-    for human_msg, ai_msg in history:
-        messages.append({"role": "user", "content": human_msg})
-        if ai_msg is not None: # ai_msg could be None if it's the first turn and history is just the user prompt
-            messages.append({"role": "assistant", "content": ai_msg})
-    messages.append({"role": "user", "content": message})
 
     print("\n--- Input to Model ---")
     print(f"System Prompt: {system_prompt if system_prompt and system_prompt.strip() else 'None'}")
-    print(f"History: {history}")
+    print(f"History (tuples format for predict): {history}")
     print(f"Current Message: {message}")
-    print(f"Formatted messages for create_chat_completion: {messages}")
+    print(f"Formatted messages for create_chat_completion: {messages_for_api}")
     print("--- End Input to Model ---\n")
 
     assistant_response_text = ""
@@ -123,14 +128,13 @@ def predict(message, history, system_prompt, max_new_tokens, temperature, top_p,
     try:
         print("Attempting generation with llm.create_chat_completion()...")
         response = llm.create_chat_completion(
-            messages=messages,
+            messages=messages_for_api,
             temperature=temperature,
             top_p=top_p,
             top_k=top_k,
             repeat_penalty=repeat_penalty,
             max_tokens=max_new_tokens,
             stop=stop_tokens,
-            # stream=True # For streaming output, Gradio handles this differently
         )
         assistant_response_text = response['choices'][0]['message']['content'].strip()
         print(f"create_chat_completion successful. Raw response: {response['choices'][0]['message']}")
@@ -139,16 +143,15 @@ def predict(message, history, system_prompt, max_new_tokens, temperature, top_p,
         print(f"Error during create_chat_completion: {e_chat_completion}")
         print("Falling back to manual prompt construction and llm()...")
 
-        # Construct prompt manually as a fallback (simplified Qwen format)
         prompt = ""
         if system_prompt and system_prompt.strip():
-            prompt += f"{im_start_token}system\n{system_prompt.strip()}{im_end_token}\n"
+            prompt += f"{IM_START_TOKEN}system\n{system_prompt.strip()}{IM_END_TOKEN}\n"
 
-        for human_msg, ai_msg in history:
-            prompt += f"{im_start_token}user\n{human_msg}{im_end_token}\n"
+        for human_msg, ai_msg in history: # history here is history_for_predict
+            prompt += f"{IM_START_TOKEN}user\n{human_msg}{IM_END_TOKEN}\n"
             if ai_msg is not None:
-                 prompt += f"{im_start_token}assistant\n{ai_msg}{im_end_token}\n"
-        prompt += f"{im_start_token}user\n{message}{im_end_token}\n{im_start_token}assistant\n" # Model should continue from here
+                 prompt += f"{IM_START_TOKEN}assistant\n{ai_msg}{IM_END_TOKEN}\n"
+        prompt += f"{IM_START_TOKEN}user\n{message}{IM_END_TOKEN}\n{IM_START_TOKEN}assistant\n"
 
         print(f"Fallback prompt: {prompt}")
 
@@ -161,7 +164,7 @@ def predict(message, history, system_prompt, max_new_tokens, temperature, top_p,
                 top_k=top_k,
                 repeat_penalty=repeat_penalty,
                 stop=stop_tokens,
-                echo=False # Don't echo the input prompt
+                echo=False
             )
             assistant_response_text = output['choices'][0]['text'].strip()
             print(f"Fallback llm() successful. Raw output: {output['choices'][0]['text']}")
@@ -190,8 +193,9 @@ def create_gradio_interface():
                     [],
                     elem_id="chatbot",
                     label="Chat Window",
-                    bubble_full_width=False,
+                    # bubble_full_width=False, # Deprecated
                     height=500,
+                    type="messages" # Use OpenAI-style messages format
                 )
                 user_input = gr.Textbox(
                     show_label=False,
@@ -208,7 +212,7 @@ def create_gradio_interface():
                     lines=3
                 )
                 max_new_tokens_slider = gr.Slider(
-                    minimum=32, maximum=N_CTX, value=DEFAULT_MAX_NEW_TOKENS, step=32, # Max tokens cannot exceed context
+                    minimum=32, maximum=N_CTX, value=DEFAULT_MAX_NEW_TOKENS, step=32,
                     label="Max New Tokens"
                 )
                 temperature_slider = gr.Slider(
@@ -227,88 +231,66 @@ def create_gradio_interface():
                     minimum=1.0, maximum=2.0, value=DEFAULT_REPEAT_PENALTY, step=0.05,
                     label="Repeat Penalty"
                 )
-                # Hidden status textbox for errors
                 status_display = gr.Textbox(label="Status", interactive=False, visible=False)
 
 
-        # Chat submission logic
-        def handle_submit(message, chat_history, sys_prompt, max_tokens, temp, top_p_val, top_k_val, rep_penalty):
-            if llm is None:
-                # Update status display if model not loaded
-                # This part is tricky as Gradio submit doesn't easily update arbitrary components outside its output
-                # For now, errors from predict will be returned in the chat.
-                # A more robust way would be a global status or specific UI element.
-                print("Attempted to chat but LLM is not loaded.")
-                # A simple way to indicate an issue if llm is None
-                chat_history.append((message, "ERROR: Model not loaded. Please check server logs."))
-                return "", chat_history, "ERROR: Model not loaded."
-
-            # Append user message
-            chat_history.append((message, None))
-            # We pass the full system prompt and params to predict
-            return "", chat_history, sys_prompt, max_tokens, temp, top_p_val, top_k_val, rep_penalty
-
-
-        # Connect user input to the generation
-        submit_args = {
-            "fn": predict,
-            "inputs": [user_input, chatbot, system_prompt_input, max_new_tokens_slider, temperature_slider, top_p_slider, top_k_slider, repeat_penalty_slider],
-            "outputs": [chatbot], # Predict will update the last AI message in chatbot
-        }
-
-        # Gradio's ChatInterface simplifies history management, but for custom layouts, we manage it manually.
-        # Here, we'll use a more direct approach like gr.Interface or manual updates.
-        # Since we use gr.Chatbot and manage history, we need to ensure `predict` gets the right state.
-        # `predict` directly takes history and returns the new AI response.
-        # Gradio's `gr.Chatbot` will automatically append the (user, ai_response) pair.
+        def user_chat_fn(user_message, chat_history_messages, sys_prompt, max_tok, temp, top_p_val, top_k_val, rep_pen):
+            if not user_message.strip(): # Do nothing if user message is empty
+                return "", chat_history_messages, sys_prompt, max_tok, temp, top_p_val, top_k_val, rep_pen
 
-        user_input.submit(
-            predict,
-            [user_input, chatbot, system_prompt_input, max_new_tokens_slider, temperature_slider, top_p_slider, top_k_slider, repeat_penalty_slider],
-            [user_input, chatbot], # Clear user_input, update chatbot
-            # The `predict` function returns only the assistant's response string.
-            # Gradio Chatbot expects the new AI message to be the output to update the last turn.
-            # So, we need a wrapper if we want to clear user_input and update chatbot
-        )
-
-        # A slightly cleaner way to handle chatbot updates with custom parameters
-        # and clearing input box:
-        def user_chat_fn(user_message, chat_history, sys_prompt, max_tok, temp, top_p_val, top_k_val, rep_pen):
             if llm is None:
-                chat_history.append((user_message, "ERROR: Model not loaded. Check logs."))
-                return "", chat_history # Clear input, update history
+                chat_history_messages.append({"role": "user", "content": user_message})
+                chat_history_messages.append({"role": "assistant", "content": "ERROR: Model not loaded. Check server logs."})
+                return "", chat_history_messages, sys_prompt, max_tok, temp, top_p_val, top_k_val, rep_pen
+
+            chat_history_messages.append({"role": "user", "content": user_message})
+            # Add a placeholder for assistant message that bot_response_fn will fill
+            chat_history_messages.append({"role": "assistant", "content": None})
+            return "", chat_history_messages, sys_prompt, max_tok, temp, top_p_val, top_k_val, rep_pen
+
+        def bot_response_fn(chat_history_messages, sys_prompt, max_tok, temp, top_p_val, top_k_val, rep_pen):
+            if llm is None or chat_history_messages[-1]["content"] is not None : # If model not loaded or already processed
+                return chat_history_messages
+
+            user_message = chat_history_messages[-2]["content"] # Get the last user message
+
+            # Convert OpenAI-style message history (List[Dict]) to tuple-style for predict's current internal logic
+            history_for_predict_tuples = []
+            # Iterate up to the second to last message (the current user's message)
+            # Each pair of (user, assistant) forms one turn for the tuple history
+            i = 0
+            temp_history = chat_history_messages[:-2] # Exclude current user and assistant placeholder
             
-            # Append user message, AI response will be None initially
-            chat_history.append((user_message, None))
-            return "", chat_history, sys_prompt, max_tok, temp, top_p_val, top_k_val, rep_pen
+            # Skip system prompt if present at the beginning for tuple conversion
+            start_index = 0
+            if temp_history and temp_history[0]["role"] == "system":
+                start_index = 1 # System prompt handled separately in predict
 
-        def bot_response_fn(chat_history, sys_prompt, max_tok, temp, top_p_val, top_k_val, rep_pen):
-            if llm is None: # Should be caught by user_chat_fn, but double check
-                return chat_history # No change
+            for i in range(start_index, len(temp_history), 2):
+                if i + 1 < len(temp_history) and temp_history[i]["role"] == "user" and temp_history[i+1]["role"] == "assistant":
+                    history_for_predict_tuples.append(
+                        (temp_history[i]["content"], temp_history[i+1]["content"])
+                    )
+                elif temp_history[i]["role"] == "user": # Handle case where last turn was only a user message (shouldn't happen if paired)
+                     history_for_predict_tuples.append((temp_history[i]["content"], None))
 
-            # The last message in history is the user's current message
-            user_message = chat_history[-1][0]
-            # The history to pass to `predict` should not include the current user turn's empty AI response
-            history_for_predict = chat_history[:-1]
 
-            bot_msg = predict(user_message, history_for_predict, sys_prompt, max_tok, temp, top_p_val, top_k_val, rep_pen)
-            chat_history[-1] = (user_message, bot_msg) # Update the last turn with AI's response
-            return chat_history
+            bot_msg_content = predict(user_message, history_for_predict_tuples, sys_prompt, max_tok, temp, top_p_val, top_k_val, rep_pen)
+            chat_history_messages[-1]["content"] = bot_msg_content # Update the assistant's placeholder message
+            return chat_history_messages
 
-        # Chain the actions: user input -> update chatbot (user msg) -> bot generates -> update chatbot (bot msg)
         user_input.submit(
             user_chat_fn,
             [user_input, chatbot, system_prompt_input, max_new_tokens_slider, temperature_slider, top_p_slider, top_k_slider, repeat_penalty_slider],
-            [user_input, chatbot, system_prompt_input, max_new_tokens_slider, temperature_slider, top_p_slider, top_k_slider, repeat_penalty_slider], # Outputs for user_chat_fn
-            queue=False # User input should be fast
+            [user_input, chatbot, system_prompt_input, max_new_tokens_slider, temperature_slider, top_p_slider, top_k_slider, repeat_penalty_slider],
+            queue=False
         ).then(
             bot_response_fn,
             [chatbot, system_prompt_input, max_new_tokens_slider, temperature_slider, top_p_slider, top_k_slider, repeat_penalty_slider],
-            [chatbot], # Output for bot_response_fn
-            queue=True # Generation can take time
+            [chatbot],
+            queue=True
         )
 
-
         gr.Examples(
             examples=[
                 ["Hello, how are you today?", "You are a friendly and helpful AI assistant specializing in concise answers."],
@@ -317,21 +299,18 @@ def create_gradio_interface():
                 ["Explain the concept of black holes to a 5-year-old.", "Keep it simple and use an analogy."]
             ],
             inputs=[user_input, system_prompt_input],
-            # outputs=[chatbot], # Examples don't directly feed to chatbot output here with this setup
-            # fn=lambda q, s: (None, [(q, predict(q, [], s, ...default_params...))]) # Complex to run predict for examples
-            # For simplicity, examples just populate the input fields.
         )
 
         with gr.Accordion("Advanced/Debug Info", open=False):
+            # Accessing global IM_START_TOKEN and IM_END_TOKEN
             gr.Markdown(f"""
             - **Model File:** `{LOCAL_MODEL_PATH}`
             - **N_CTX:** `{N_CTX}`
             - **N_THREADS:** `{N_THREADS if N_THREADS is not None else 'Auto'}`
             - **N_GPU_LAYERS:** `{N_GPU_LAYERS}`
             - **Log Verbosity (llama.cpp):** `{VERBOSE_LLAMA}`
-            - **Stop Tokens Used:** `{im_start_token}system`, `{im_start_token}user`, `{im_end_token}`, `EOS_TOKEN`
+            - **Stop Tokens Used (Conceptual):** `{IM_START_TOKEN}system`, `{IM_START_TOKEN}user`, `{IM_END_TOKEN}`, `EOS_TOKEN`
             """)
-            # Add a button to attempt model reload if it failed initially
             reload_button = gr.Button("Attempt to Reload Model")
             reload_status = gr.Label(value="Model Status: Unknown")
 
@@ -343,7 +322,17 @@ def create_gradio_interface():
 
             def attempt_reload():
                 global llm
-                llm = None # Force re-evaluation of loading
+                if llm is not None:
+                    try:
+                        # Attempt to free existing model if Llama.cpp supports it or by reassigning
+                        print("Freeing existing model instance (if any)...")
+                        del llm # Explicitly delete to trigger __del__ if possible
+                        llm = None
+                        import gc
+                        gc.collect() # Suggest garbage collection
+                    except Exception as e_del:
+                        print(f"Error during manual deletion of llm: {e_del}")
+                
                 if load_llm_model():
                     return "Model reloaded successfully!"
                 else:
@@ -351,8 +340,6 @@ def create_gradio_interface():
 
             reload_button.click(attempt_reload, outputs=[reload_status])
             iface.load(update_reload_status, outputs=[reload_status]) # Update status on interface load
-
-
     return iface
 
 # --- Main Execution ---
@@ -362,19 +349,15 @@ if __name__ == "__main__":
 
     if model_available:
         if not load_llm_model():
-            print("Model loading failed. The Gradio interface will start, but chat functionality will be impaired.")
-            print("You can try to reload the model via the 'Advanced/Debug Info' section in the UI.")
+            print("Initial model loading failed. Gradio will start; use UI to attempt reload.")
         else:
             print("Model ready.")
     else:
-        print("Model download failed. Cannot proceed to load model or start chat functionality.")
-        print("The Gradio interface will start, but it will not be functional.")
+        print("Model download failed. Cannot load model. Gradio will start; chat will be non-functional.")
 
     print("Creating Gradio interface...")
     app_interface = create_gradio_interface()
 
     print("Launching Gradio interface...")
-    # Share=True is useful for public links if running locally, but HF Spaces handles public URL.
-    # In_browser=True to open in browser locally.
     app_interface.launch()
-    print("Gradio interface launched. Check your terminal or logs for the URL.")
+    print("Gradio interface launched.")