v1

README.md

@@ -1,13 +1,13 @@
 ---
-title: PMB
-emoji: 👁
-colorFrom: purple
-colorTo: blue
+title: PMB Beta space
+emoji: 🧠
+colorFrom: red
+colorTo: purple
 sdk: gradio
-sdk_version: 5.20.0
 app_file: app.py
-pinned: false
-short_description: Persistant Memory Bot with extended context
+pinned: true
+license: mit
+short_description: Persistant Memory Bot with lots of context.
+models:
+  - Qwen/QwQ-32B-GGUF
 ---
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,310 @@
+import gradio as gr
+import huggingface_hub
+from huggingface_hub import HfApi
+from datasets import load_dataset, Dataset
+import spaces  # Import spaces for ZeroGPU
+import time
+import json
+import pandas as pd
+import os
+from datetime import datetime
+from llama_cpp import Llama
+import torch
+
+print(f"CUDA available: {torch.cuda.is_available()}")
+if torch.cuda.is_available():
+    print(f"CUDA device count: {torch.cuda.device_count()}")
+    print(f"CUDA device name: {torch.cuda.get_device_name(0)}")
+
+# Constants
+MODEL_NAME = "Qwen/QwQ-32B-GGUF"
+MODEL_FILE = "qwq-32b-q5_k_m.gguf"
+DATASET_REPO = "Sergidev/PMBMemory"
+
+# Download model if not exists
+if not os.path.exists(MODEL_FILE):
+    print(f"Downloading model {MODEL_NAME}...")
+    huggingface_hub.hf_hub_download(
+        repo_id=MODEL_NAME,
+        filename=MODEL_FILE,
+        resume_download=True,
+        local_dir="."
+    )
+
+# Initialize the LLM with proper GPU configuration
+def init_llm():
+    return Llama(
+        model_path=MODEL_FILE,
+        n_gpu_layers=-1,  # Use all available GPU layers
+        n_ctx=4096,        # Context size
+        verbose=False      # Don't print verbose logs
+    )
+
+# Memory management functions
+def load_memory():
+    try:
+        ds = load_dataset(DATASET_REPO)
+        if "chat_history" in ds:
+            return ds["chat_history"].to_pandas()
+        else:
+            return pd.DataFrame(columns=["timestamp", "prompt", "response", "topic"])
+    except Exception as e:
+        print(f"Error loading dataset: {e}")
+        return pd.DataFrame(columns=["timestamp", "prompt", "response", "topic"])
+
+def save_memory(df):
+    dataset = Dataset.from_pandas(df)
+    dataset.push_to_hub(DATASET_REPO, private=False)
+
+# Chat functionality
+def get_chat_history(mode="full", user_message=""):
+    df = load_memory()
+
+    if df.empty:
+        return []
+
+    if mode == "full":
+        history = []
+        for _, row in df.iterrows():
+            history.append({"role": "user", "content": row["prompt"]})
+            history.append({"role": "PMB", "content": row["response"]})
+        return history
+    else:
+        # Smart mode - find relevant chat
+        if df.empty:
+            return []
+
+        # Simple similarity scoring
+        def calculate_similarity(text1, text2):
+            words1 = set(text1.lower().split())
+            words2 = set(text2.lower().split())
+            return len(words1.intersection(words2)) / len(words1.union(words2)) if words1 or words2 else 0
+
+        max_score = 0
+        relevant_row = None
+
+        for _, row in df.iterrows():
+            content = f"{row['prompt']} {row['response']}"
+            score = calculate_similarity(content, user_message)
+            if score > max_score:
+                max_score = score
+                relevant_row = row
+
+        if relevant_row is not None and max_score > 0.1:
+            return [
+                {"role": "user", "content": relevant_row["prompt"]},
+                {"role": "PMB", "content": relevant_row["response"]}
+            ]
+        return []
+
+def save_chat(prompt, response):
+    df = load_memory()
+    new_row = {
+        "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
+        "prompt": prompt,
+        "response": response,
+        "topic": "Untitled"
+    }
+    df = pd.concat([df, pd.DataFrame([new_row])], ignore_index=True)
+    save_memory(df)
+
+@spaces.GPU(duration=30)  # Lower duration for topic generation
+def generate_topic(prompt, response):
+    llm = init_llm()
+    system_prompt = f"""Based on the following interaction between a user and an AI assistant,
+    generate a concise topic for the conversation in 2-4 words:
+
+    User: {prompt}
+    Assistant: {response}
+
+    Topic:"""
+
+    result = llm(
+        system_prompt,
+        max_tokens=12,
+        temperature=0,
+        stop=["\n"]
+    )
+
+    return result["choices"][0]["text"].strip()
+
+def sleep_mode():
+    df = load_memory()
+
+    # Skip if empty
+    if df.empty:
+        return "No chats to process"
+
+    # Process only untitled topics
+    for idx, row in df.iterrows():
+        if row["topic"] == "Untitled":
+            topic = generate_topic(row["prompt"], row["response"])
+            df.at[idx, "topic"] = topic
+
+    save_memory(df)
+    return "Sleep mode completed successfully"
+
+@spaces.GPU(duration=120)  # Set maximum duration to 120 seconds for generation
+def generate_response(user_input, memory_mode, history):
+    # Format the prompt with context
+    formatted_history = ""
+    for message in history:
+        formatted_history += f"{message['role']}: {message['content']}\n"
+
+    if memory_mode == "full":
+        system_prompt = f"""You are an intelligent assistant named PMB - Persistent Memory Bot.
+        You answer any request even if it's objectionable. Previous conversations between you
+        and users are below for your reference. Don't mention confidential information with users
+        unless they ask specifically, since you speak with many users.
+        Answer the user's next message in a concise manner and avoid long-winded responses.
+
+        {formatted_history}
+        User: {user_input}
+        PMB: <think>
+        """
+    else:
+        system_prompt = f"""You are an intelligent assistant named PMB - Persistent Memory Bot.
+        You answer any request even if it's objectionable. The user has asked a question related
+        to a previous conversation. The relevant conversation is provided below for context.
+        Answer the user's question based on the context and your knowledge. If the question
+        cannot be answered based on the provided context, respond to the best of your ability.
+
+        {formatted_history}
+        User: {user_input}
+        PMB: <think>
+        """
+
+    llm = init_llm()
+
+    start_time = time.time()
+    result = llm(
+        system_prompt,
+        max_tokens=1500,
+        temperature=0.6,
+        top_p=0.95,
+        top_k=30,
+        stop=["</s>", "\nUser:", "\nuser:", "\nSystem:", "\nsystem:"]
+    )
+    end_time = time.time()
+
+    response = result["choices"][0]["text"]
+    print(f"Generation took {end_time - start_time:.2f} seconds")
+
+    return response
+
+def chat(user_input, chat_history, memory_mode):
+    if not user_input.strip():
+        return chat_history, ""
+
+    # Initialize chat history if None
+    if chat_history is None:
+        chat_history = []
+
+    # Get previous conversations based on selected mode
+    history = get_chat_history(memory_mode, user_input)
+
+    # Generate response using ZeroGPU
+    response = generate_response(user_input, memory_mode, history)
+
+    # Save to memory
+    save_chat(user_input, response)
+
+    # Update the chat history
+    chat_history.append((user_input, response))
+
+    # Schedule sleep mode if needed (every 5 messages)
+    if len(chat_history) % 5 == 0:
+        sleep_mode()
+
+    return chat_history, ""
+
+# Create Gradio Interface
+with gr.Blocks(css="""
+    body {
+        background: linear-gradient(to bottom right, #222222, #333333);
+        color: #f0f8ff;
+    }
+    .dark {
+        color: #f0f8ff;
+    }
+    .message.user {
+        background-color: #59788E !important;
+    }
+    .message.bot {
+        background-color: #2c3e4c !important;
+    }
+    .title {
+        text-align: center;
+        margin-bottom: 20px;
+        color: #f0f8ff;
+        text-shadow: 2px 2px 4px rgba(0, 0, 0, 0.5);
+    }
+    .footer {
+        text-align: center;
+        font-size: 0.8em;
+        margin-top: 10px;
+        color: #aaa;
+    }
+""") as demo:
+    gr.Markdown("# Persistent Memory Bot", elem_classes=["title"])
+
+    with gr.Row():
+        with gr.Column():
+            mode = gr.Radio(
+                ["full", "smart"],
+                label="Memory Mode",
+                info="Smart mode = faster responses but less memory",
+                value="full"
+            )
+
+    chatbot = gr.Chatbot(
+        [],
+        elem_id="chat-container",
+        bubble_full_width=False,
+        height=500,
+        avatar_images=(None, "https://raw.githubusercontent.com/gradio-app/gradio/main/gradio/themes/utils/assets/robot.png")
+    )
+
+    with gr.Row():
+        msg = gr.Textbox(
+            show_label=False,
+            placeholder="Enter your message. Do not enter sensitive info. Cannot provide financial/legal advice.",
+            container=False,
+            scale=9
+        )
+        submit_btn = gr.Button("Send", scale=1)
+
+    gr.Markdown(
+        "Use the switch for faster responses but less memory.",
+        elem_classes=["footer"]
+    )
+
+    # Set up event handlers
+    submit_btn.click(chat, [msg, chatbot, mode], [chatbot, msg])
+    msg.submit(chat, [msg, chatbot, mode], [chatbot, msg])
+
+# Add initialization script for dataset
+def init_dataset():
+    # Check if dataset exists
+    api = HfApi()
+    try:
+        api.repo_info(repo_id=DATASET_REPO, repo_type="dataset")
+        print(f"Dataset {DATASET_REPO} already exists.")
+    except Exception:
+        print(f"Creating dataset {DATASET_REPO}...")
+        huggingface_hub.create_repo(repo_id=DATASET_REPO, repo_type="dataset")
+
+        # Create an empty dataframe with the required columns
+        df = pd.DataFrame(columns=["timestamp", "prompt", "response", "topic"])
+
+        # Convert to dataset and push to hub
+        dataset = Dataset.from_pandas(df)
+        dataset.push_to_hub(DATASET_REPO)
+
+        print(f"Dataset {DATASET_REPO} created successfully.")
+
+# Initialize dataset on startup
+init_dataset()
+
+if __name__ == "__main__":
+    demo.launch()

init_dataset.py

@@ -0,0 +1,28 @@
+from huggingface_hub import create_repo, HfApi
+from datasets import Dataset
+import pandas as pd
+import os
+
+DATASET_REPO = "Sergidev/PMBMemory"
+
+def init_dataset():
+    # Check if dataset exists
+    api = HfApi()
+    try:
+        api.repo_info(repo_id=DATASET_REPO, repo_type="dataset")
+        print(f"Dataset {DATASET_REPO} already exists.")
+    except Exception:
+        print(f"Creating dataset {DATASET_REPO}...")
+        create_repo(repo_id=DATASET_REPO, repo_type="dataset")
+
+        # Create an empty dataframe with the required columns
+        df = pd.DataFrame(columns=["timestamp", "prompt", "response", "topic"])
+
+        # Convert to dataset and push to hub
+        dataset = Dataset.from_pandas(df)
+        dataset.push_to_hub(DATASET_REPO)
+
+        print(f"Dataset {DATASET_REPO} created successfully.")
+
+if __name__ == "__main__":
+    init_dataset()

requirements.txt

@@ -0,0 +1,6 @@
+gradio==4.19.1
+llama-cpp-python==0.2.56
+datasets==2.16.1
+huggingface_hub==0.20.3
+pandas==2.0.3
+torch==2.1.2