trying v3 fix app.py

app.py

@@ -8,14 +8,17 @@ import uuid
 import json
 import re
 import random
-from typing import List, Tuple, Dict, Any
 import hashlib
+from typing import List, Tuple, Dict, Any
 
-# Set device globally
+# ------------------------------
+# ⚙️ Global Config
+# ------------------------------
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+MODEL_NAME = "bhushanrocks/supportpal-dialoGPT-v3"
 
 # ------------------------------
-# 🔐 MongoDB Setup & Auth Management
+# 🔐 MongoDB Setup & Local Fallback
 # ------------------------------
 MONGO_URI = os.getenv("MONGO_URI")
 use_mongo = False
@@ -33,395 +36,278 @@ if MONGO_URI:
         use_mongo = True
         print("✅ Connected to MongoDB for persistence and authentication.")
     except Exception as e:
-        # NOTE: In a Hugging Face Space, network connections can sometimes fail. 
-        # Falling back to local storage is the correct behavior.
         print(f"⚠️ MongoDB unavailable: {e}. Using local JSON storage.")
 else:
     print("⚠️ MONGO_URI not set. Using local JSON storage.")
 
-# ------------------------------
-# 💾 Local JSON Fallback (User Auth and Chat History)
-# ------------------------------
 LOCAL_USERS_FILE = "local_users.json"
 LOCAL_CHATS_FILE = "local_chats.json"
 
+
 def load_local_data(filepath):
     if os.path.exists(filepath):
         with open(filepath, "r") as f:
             return json.load(f)
     return {}
 
+
 def save_local_data(data, filepath):
     with open(filepath, "w") as f:
         json.dump(data, f, indent=4)
 
+
 local_users = load_local_data(LOCAL_USERS_FILE)
 local_chats = load_local_data(LOCAL_CHATS_FILE)
 
-# Utility to hash passwords for storage (even locally)
+
 def hash_password(password):
     return hashlib.sha256(password.encode()).hexdigest()
 
 # ------------------------------
-# 🔑 Authentication Functions
+# 🔑 Auth Functions
 # ------------------------------
-
 def authenticate_user(username, password):
-    hashed_password = hash_password(password)
-    
+    hashed = hash_password(password)
     if use_mongo:
         user_data = users_collection.find_one({"username": username})
     else:
         user_data = local_users.get(username)
 
-    if user_data and user_data['password'] == hashed_password:
-        return user_data['user_id'], f"Welcome back, {username}!"
+    if user_data and user_data["password"] == hashed:
+        return user_data["user_id"], f"Welcome back, {username}!"
     return None, "Invalid username or password."
 
+
 def register_user(username, password):
-    hashed_password = hash_password(password)
+    hashed = hash_password(password)
     new_user_id = str(uuid.uuid4())
-    
+
     if use_mongo:
         if users_collection.find_one({"username": username}):
             return None, "Username already taken."
-        users_collection.insert_one({"username": username, "password": hashed_password, "user_id": new_user_id})
+        users_collection.insert_one({"username": username, "password": hashed, "user_id": new_user_id})
     else:
         if username in local_users:
             return None, "Username already taken."
-        local_users[username] = {"username": username, "password": hashed_password, "user_id": new_user_id}
+        local_users[username] = {"username": username, "password": hashed, "user_id": new_user_id}
         save_local_data(local_users, LOCAL_USERS_FILE)
-        
+
     return new_user_id, f"Account created! Welcome, {username}."
 
 # ------------------------------
-# 📜 Conversation Management & Gradio Format Conversion
+# 💾 Chat History Handling
 # ------------------------------
-
-# Helper to convert stored tuple history to Gradio 'messages' format
-def history_to_gradio_format(history: List[Tuple[str, str]]) -> List[Dict[str, str]]:
-    """Converts [(user_msg, bot_msg)] tuples to Gradio message dictionaries."""
-    gradio_history = []
-    for user_msg, bot_msg in history:
-        if user_msg:
-            gradio_history.append({"role": "user", "content": user_msg})
-        if bot_msg:
-            gradio_history.append({"role": "assistant", "content": bot_msg})
-    return gradio_history
-
 def get_history(user_id):
-    """Fetch user conversation history (stored as tuples)."""
     if not user_id:
         return []
-    
     if use_mongo:
         convo = chats_collection.find_one({"user_id": user_id})
-        # History is stored as list of tuples: [(user_msg, bot_msg), ...]
         return convo.get("history", [])
-    else:
-        return local_chats.get(user_id, [])
+    return local_chats.get(user_id, [])
+
 
-def save_history(user_id, history: List[Tuple[str, str]]):
-    """Save conversation history (stored as tuples)."""
+def save_history(user_id, history):
     if not user_id:
         return
-    
     if use_mongo:
-        chats_collection.update_one(
-            {"user_id": user_id}, 
-            {"$set": {"history": history}}, 
-            upsert=True
-        )
+        chats_collection.update_one({"user_id": user_id}, {"$set": {"history": history}}, upsert=True)
     else:
         local_chats[user_id] = history
         save_local_data(local_chats, LOCAL_CHATS_FILE)
 
+
+def history_to_gradio_format(history: List[Tuple[str, str]]):
+    gradio_msgs = []
+    for user_msg, bot_msg in history:
+        if user_msg:
+            gradio_msgs.append({"role": "user", "content": user_msg})
+        if bot_msg:
+            gradio_msgs.append({"role": "assistant", "content": bot_msg})
+    return gradio_msgs
+
 # ------------------------------
-# 🤗 Load Model (PEFT/LoRA Integration)
+# 🤗 Load Model
 # ------------------------------
-# Assuming the user's latest checkpoint is pushed here.
-MODEL_NAME = "bhushanrocks/supportpal-dialoGPT-v3" 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 
 try:
-    # 1. Load the PEFT config from the Hub to get the base model name
     config = PeftConfig.from_pretrained(MODEL_NAME)
-    base_model_name = config.base_model_name_or_path
-    
-    # 2. Load the base model (e.g., microsoft/DialoGPT-medium)
     base_model = AutoModelForCausalLM.from_pretrained(
-        base_model_name,
-        torch_dtype=torch.float16 if DEVICE == 'cuda' else None,
+        config.base_model_name_or_path,
+        torch_dtype=torch.float16 if DEVICE == "cuda" else None,
     ).to(DEVICE)
-    
-    # 3. Load the LoRA adapters and attach them to the base model
     model = PeftModel.from_pretrained(base_model, MODEL_NAME)
-    print(f"✅ Successfully loaded PEFT model from {MODEL_NAME}")
-    
+    print(f"✅ Loaded PEFT model from {MODEL_NAME}")
 except Exception as e:
-    # Fallback if the PEFT load fails (e.g., if it was already merged or structure changed)
-    print(f"⚠️ PEFT model loading failed. Falling back to AutoModelForCausalLM. Error: {e}")
+    print(f"⚠️ Failed to load PEFT adapters, fallback: {e}")
     model = AutoModelForCausalLM.from_pretrained(MODEL_NAME).to(DEVICE)
 
 model.eval()
 
 # ------------------------------
-# 💬 Chat Logic (Refined Context Handling)
+# 💬 Chat Logic
 # ------------------------------
+def generate_response(user_input: str, state_dict: dict):
+    """Main generation logic."""
+    if not state_dict.get("logged_in"):
+        return [], state_dict, "Please log in to start chatting."
+
+    user_id = state_dict["user_id"]
+    history = state_dict.get("history", [])
+
+    # Define system prompt here
+    system_prompt = (
+        "You are SupportPal — a kind, compassionate AI that listens patiently, "
+        "validates emotions, and gently helps people feel understood. "
+        "Avoid talking about yourself, keep responses concise, warm, and natural."
+    )
 
-def generate_response(user_input, chat_history, model, tokenizer, system_prompt, max_new_tokens=150):
-    """
-    Generates a model response using the trained DialoGPT with robust handling
-    for missing AI markers, truncation, and empty outputs.
-    """
+    # Keep only the last 5 turns
+    context_turns = history[-5:]
 
-    # ------------------------------------------------------------------
-    # 1️⃣  Prepare dialogue context (robust handling)
-    # ------------------------------------------------------------------
+    # Build dialogue text
     prompt_turns = []
-
-    # Handle both dict-based and tuple/string-based history
-    for turn in chat_history:
-        if isinstance(turn, dict) and "user" in turn and "bot" in turn:
-            # Expected structured history
-            prompt_turns.append(f"Human: {turn['user']}")
-            prompt_turns.append(f"AI: {turn['bot']}")
+    for turn in context_turns:
+        if isinstance(turn, dict):
+            prompt_turns.append(f"Human: {turn.get('user','')}\nAI: {turn.get('bot','')}")
         elif isinstance(turn, (list, tuple)) and len(turn) == 2:
-            # Gradio ChatInterface sometimes stores (user, bot) tuples
-            prompt_turns.append(f"Human: {turn[0]}")
-            prompt_turns.append(f"AI: {turn[1]}")
-        elif isinstance(turn, str):
-            # Fallback: unexpected raw string history
-            prompt_turns.append(f"Human: {turn}")
+            prompt_turns.append(f"Human: {turn[0]}\nAI: {turn[1]}")
 
     dialogue = system_prompt.strip() + "\n\n" + "\n".join(prompt_turns)
-    dialogue += f"\nHuman: {user_input.strip()}\nAI:"
-
-    # ------------------------------------------------------------------
-    # 2️⃣  Tokenize safely (preserve recent turns)
-    # ------------------------------------------------------------------
-    tokenizer.truncation_side = "left"  # keep last turns
-    inputs = tokenizer(
-        dialogue + tokenizer.eos_token,
-        return_tensors="pt",
-        truncation=True,
-        max_length=512,
-    ).to(model.device)
-
-    # ------------------------------------------------------------------
-    # 3️⃣  Generate
-    # ------------------------------------------------------------------
+    dialogue += f"\nHuman: {user_input}\nAI:"
+
+    # Tokenize and generate
+    inputs = tokenizer(dialogue, return_tensors="pt", truncation=True, max_length=512).to(model.device)
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
-            max_new_tokens=max_new_tokens,
-            temperature=0.8,
+            max_new_tokens=200,
+            temperature=0.7,
+            top_k=50,
             top_p=0.9,
-            repetition_penalty=1.1,
             do_sample=True,
             pad_token_id=tokenizer.eos_token_id,
         )
 
     decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    response = decoded.split("AI:")[-1].strip()
+
+    # Clean and finalize
+    response = re.sub(r"You are SupportPal.*", "", response, flags=re.DOTALL).strip()
+    response = re.sub(r"\b(I am|I'm) glad\b.*?\.", "", response)
+    response = re.sub(r"\b(I hope|I wish)\b.*?\.", "", response)
+    response = response.replace("AI:", "").strip()
 
-    # ------------------------------------------------------------------
-    # 4️⃣  Extract response robustly
-    # ------------------------------------------------------------------
-    # Look for the last "AI:" marker and take everything after it
-    match = re.search(r"AI:\s*(.*)", decoded, flags=re.DOTALL)
-    response = match.group(1).strip() if match else decoded.strip()
+    if not response:
+        response = "I'm here with you. Tell me more about what’s on your mind. 💛"
 
-    # Stop before the next Human turn if model over-generated
-    response = re.split(r"\nHuman:", response)[0].strip()
+    if random.random() < 0.25 and not any(x in response for x in ["💛", "✨"]):
+        response += " " + random.choice(["💛", "✨", "You’re not alone in this."])
 
-    # Clean out any duplicated system prompt or extra markers
-    response = re.sub(rf"{re.escape(system_prompt)}", "", response).strip()
-    response = re.sub(r"AI:\s*", "", response).strip()
+    # Save and return
+    history.append((user_input, response))
+    save_history(user_id, history)
+    state_dict["history"] = history
 
-    # ------------------------------------------------------------------
-    # 5️⃣  Handle empty or invalid responses
-    # ------------------------------------------------------------------
-    if not response or response.lower() in ["", "human:", "ai:"]:
-        response = "I'm here with you. Could you tell me a bit more about how you're feeling?"
+    gradio_history = history_to_gradio_format(history)
+    return gradio_history, state_dict, ""
 
-    # Optional empathy touch 💛
-    if random.random() < 0.25 and not response.endswith(("💛", "✨", ".", "?", "!")):
-        response += " 💛"
 
-    # ------------------------------------------------------------------
-    # 6️⃣  Update chat history & return
-    # ------------------------------------------------------------------
-    chat_history.append({"user": user_input, "bot": response})
-    return response, chat_history
+def process_chat(user_message, state_dict):
+    """Wrapper used for Gradio callbacks."""
+    if not state_dict:
+        state_dict = {"logged_in": False, "user_id": None, "username": None, "history": []}
+
+    if not state_dict.get("logged_in"):
+        return [], state_dict, "Please log in to start chatting."
+
+    return generate_response(user_message, state_dict)
 
 # ------------------------------
-# 🔑 Gradio Authentication Handlers
+# 🔐 Login / Register / Logout
 # ------------------------------
-
-def handle_login(username, password, current_state):
-    user_id, status_message = authenticate_user(username, password)
-    
+def handle_login(username, password, state):
+    user_id, status = authenticate_user(username, password)
     if user_id:
-        # Get history in tuple format
-        history = get_history(user_id) 
-        
-        current_state.update({
-            'logged_in': True, 
-            'user_id': user_id, 
-            'username': username,
-            'history': history, # Store history in tuple format
-        })
-        
-        # Convert history to Gradio 'messages' format for display
-        gradio_history = history_to_gradio_format(history)
-        
-        # Return components: status, auth_panel hide, chat_panel show, chatbot history (gradio format), state
-        return status_message, gr.update(visible=False), gr.update(visible=True), gradio_history, current_state, f"Logged in as: **{username}**"
-    else:
-        # Return components: status, auth_panel show, chat_panel hide, empty history, state, user status (no change)
-        return status_message, gr.update(visible=True), gr.update(visible=False), [], current_state, gr.update()
+        history = get_history(user_id)
+        state.update({"logged_in": True, "user_id": user_id, "username": username, "history": history})
+        gr_history = history_to_gradio_format(history)
+        return status, gr.update(visible=False), gr.update(visible=True), gr_history, state, f"Logged in as: **{username}**"
+    return status, gr.update(visible=True), gr.update(visible=False), [], state, gr.update()
+
 
-def handle_register(username, password, current_state):
-    user_id, status_message = register_user(username, password)
-    
+def handle_register(username, password, state):
+    user_id, status = register_user(username, password)
     if user_id:
-        current_state.update({
-            'logged_in': True, 
-            'user_id': user_id, 
-            'username': username,
-            'history': [], # Store history in tuple format
-        })
-        # Return components: status, auth_panel hide, chat_panel show, empty history (gradio format), state
-        return status_message, gr.update(visible=False), gr.update(visible=True), [], current_state, f"Logged in as: **{username}**"
-    else:
-        # Return components: status, auth_panel show, chat_panel hide, empty history, state, user status (no change)
-        return status_message, gr.update(visible=True), gr.update(visible=False), [], current_state, gr.update()
+        state.update({"logged_in": True, "user_id": user_id, "username": username, "history": []})
+        return status, gr.update(visible=False), gr.update(visible=True), [], state, f"Logged in as: **{username}**"
+    return status, gr.update(visible=True), gr.update(visible=False), [], state, gr.update()
+
 
 def handle_logout():
-    # Reset all components and states
     return (
-        "Logged out successfully. Please log in again.", 
-        gr.update(visible=True), # Show auth panel
-        gr.update(visible=False), # Hide chat panel
-        [], # Clear chatbot history
-        {'logged_in': False, 'user_id': None, 'username': None, 'history': []}, # Reset state
-        gr.update(value=""), # Clear user input box
-        "Logged in as:" # Reset user status label
+        "Logged out successfully.",
+        gr.update(visible=True),
+        gr.update(visible=False),
+        [],
+        {"logged_in": False, "user_id": None, "username": None, "history": []},
+        gr.update(value=""),
+        "Logged in as:",
     )
 
-def handle_clear_chat(current_state):
-    # Only clear history for the currently logged-in user, but keep the session alive
-    if current_state.get('logged_in'):
-        user_id = current_state['user_id']
-        save_history(user_id, []) # Save empty history (tuple format)
-        current_state['history'] = []
-        return [], current_state # Return empty Gradio history
-    return [], current_state
 
+def handle_clear_chat(state):
+    if state.get("logged_in"):
+        user_id = state["user_id"]
+        save_history(user_id, [])
+        state["history"] = []
+        return [], state
+    return [], state
 
 # ------------------------------
 # 🧱 Gradio UI
 # ------------------------------
-
-def process_chat(user_message, chat_history):
-    # Ensure consistent history format before sending to model
-    if chat_history is None:
-        chat_history = []
-    elif isinstance(chat_history, list) and len(chat_history) > 0 and isinstance(chat_history[0], (list, tuple)):
-        # Convert tuple-based history into dict-based
-        chat_history = [{"user": u, "bot": b} for u, b in chat_history]
-
-    # Call model generator
-    response, updated_history = generate_response(user_message, chat_history, model, tokenizer, system_prompt)
-    return response, updated_history
-
 with gr.Blocks(title="SupportPal Chatbot") as demo:
-    storage_mode = "☁️ MongoDB Connected" if use_mongo else "💾 Using Local Storage (Dev Only)"
-    
-    # State component to hold session data (logged_in status, user_id, history in tuple format)
-    session_state = gr.State(
-        value={'logged_in': False, 'user_id': None, 'username': None, 'history': []}
-    )
+    storage_mode = "☁️ MongoDB Connected" if use_mongo else "💾 Local Storage"
 
-    gr.Markdown(f"## 🤖 SupportPal - Empathetic Chatbot\n**Persistence Status:** {storage_mode}")
+    session_state = gr.State(value={"logged_in": False, "user_id": None, "username": None, "history": []})
+    gr.Markdown(f"## 🤖 SupportPal — Empathetic Chatbot\n**Persistence Mode:** {storage_mode}")
 
-    # --- Authentication Panel (Visible by default) ---
+    # Auth Panel
     with gr.Row(visible=True) as auth_panel:
-        with gr.Column(scale=1):
+        with gr.Column():
             auth_status = gr.Markdown("Please log in or register to start chatting.")
             username_box = gr.Textbox(label="Username")
             password_box = gr.Textbox(label="Password", type="password")
-            
             with gr.Row():
                 login_btn = gr.Button("Login", variant="primary")
                 register_btn = gr.Button("Register", variant="secondary")
 
-    # --- Chat Panel (Hidden by default) ---
+    # Chat Panel
     with gr.Column(visible=False) as chat_panel:
-        # TARGET FOR USER STATUS UPDATE (Issue 1 fix)
-        user_status = gr.Markdown("Logged in as:") 
-        
-        # NOTE: Added type='messages' to avoid the deprecation warning
-        chatbot_ui = gr.Chatbot(height=400, type='messages') 
-        
+        user_status = gr.Markdown("Logged in as:")
+        chatbot_ui = gr.Chatbot(height=400, type="messages")
         with gr.Row():
-            msg = gr.Textbox(label="Type your message here...", scale=4)
-            send_btn = gr.Button("Send", scale=1, variant="primary")
-
+            msg = gr.Textbox(label="Your message...", scale=4)
+            send_btn = gr.Button("Send", variant="primary", scale=1)
         with gr.Row():
-            clear_btn = gr.Button("Clear History", variant="secondary")
+            clear_btn = gr.Button("Clear Chat", variant="secondary")
             logout_btn = gr.Button("Logout", variant="stop")
 
-    # --- Authentication Flow (Issue 1 fix: add user_status output) ---
-    login_btn.click(
-        fn=handle_login, 
-        inputs=[username_box, password_box, session_state], 
-        outputs=[auth_status, auth_panel, chat_panel, chatbot_ui, session_state, user_status]
-    )
+    # Bind events
+    login_btn.click(handle_login, [username_box, password_box, session_state],
+                    [auth_status, auth_panel, chat_panel, chatbot_ui, session_state, user_status])
+    register_btn.click(handle_register, [username_box, password_box, session_state],
+                       [auth_status, auth_panel, chat_panel, chatbot_ui, session_state, user_status])
 
-    register_btn.click(
-        fn=handle_register, 
-        inputs=[username_box, password_box, session_state], 
-        outputs=[auth_status, auth_panel, chat_panel, chatbot_ui, session_state, user_status]
-    )
-    
-    # --- Chat Flow (FIXED) ---
-    
-    # 1. Action triggered by clicking the Send button
-    send_btn.click(
-        fn=generate_response, 
-        inputs=[msg, session_state], 
-        outputs=[chatbot_ui, session_state, auth_status] # chatbot_ui now expects 'messages' dict format
-    ).then(
-        # Clear input box after sending
-        fn=lambda: gr.update(value=""), 
-        inputs=None, 
-        outputs=[msg]
-    )
+    send_btn.click(process_chat, [msg, session_state],
+                   [chatbot_ui, session_state, auth_status]).then(lambda: gr.update(value=""), None, [msg])
 
-    # 2. Action triggered by hitting Enter in the Textbox (using submit, not click)
-    msg.submit(
-        fn=process_chat, 
-        inputs=[msg, session_state], 
-        outputs=[chatbot_ui, session_state, auth_status] # chatbot_ui now expects 'messages' dict format
-    ).then(
-        # Clear input box after sending
-        fn=lambda: gr.update(value=""), 
-        inputs=None, 
-        outputs=[msg]
-    )
+    msg.submit(process_chat, [msg, session_state],
+               [chatbot_ui, session_state, auth_status]).then(lambda: gr.update(value=""), None, [msg])
 
-    # Clear and Logout handlers (Issue 1 fix: add user_status output for logout)
-    clear_btn.click(
-        fn=handle_clear_chat,
-        inputs=[session_state],
-        outputs=[chatbot_ui, session_state]
-    )
-    
-    logout_btn.click(
-        fn=handle_logout,
-        inputs=None,
-        outputs=[auth_status, auth_panel, chat_panel, chatbot_ui, session_state, msg, user_status]
-    )
+    clear_btn.click(handle_clear_chat, [session_state], [chatbot_ui, session_state])
+    logout_btn.click(handle_logout, None,
+                     [auth_status, auth_panel, chat_panel, chatbot_ui, session_state, msg, user_status])
 
 demo.launch()