v1

modules/pmbl.py

@@ -1,15 +1,16 @@
 import sqlite3
 from datetime import datetime
 from llama_cpp import Llama
-from .hippocampus import generate_topic
+from concurrent.futures import ThreadPoolExecutor
 
 class PMBL:
     def __init__(self, model_path):
-        self.llm = Llama(model_path=model_path, n_ctx=13000, n_threads=8, n_gpu_layers=32)
+        self.model_path = model_path
         self.init_db()
+        self.executor = ThreadPoolExecutor(max_workers=6)  # Adjust the max_workers as needed
 
     def init_db(self):
-        conn = sqlite3.connect('../chat_history.db')
+        conn = sqlite3.connect('./chat_history.db')
         c = conn.cursor()
         c.execute('''CREATE TABLE IF NOT EXISTS chats
                      (id INTEGER PRIMARY KEY AUTOINCREMENT,
@@ -93,7 +94,15 @@ class PMBL:
         else:  # mode == "smart"
             system_prompt = f"You are an intelligent assistant named PMB - Persistent Memory Bot. You answer any 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.\n\n{formatted_history}\nPMB:"
 
-        response = self.llm(
+        response = self.executor.submit(self.generate_response_task, system_prompt, prompt)
+
+        for chunk in response.result():
+            yield chunk
+
+    def generate_response_task(self, system_prompt, prompt):
+        llm = Llama(model_path=self.model_path, n_ctx=13000, n_threads=8, n_gpu_layers=32)
+
+        response = llm(
             system_prompt,
             max_tokens=1500,
             temperature=0.7,
@@ -118,8 +127,23 @@ class PMBL:
 
         for chat in untitled_chats:
             chat_id, prompt, response = chat
-            topic = generate_topic(prompt, response)
+            topic = self.generate_topic(prompt, response)
             c.execute("UPDATE chats SET topic = ? WHERE id = ?", (topic, chat_id))
             conn.commit()
 
-        conn.close()
+        conn.close()
+
+    def generate_topic(self, prompt, response):
+        llm = Llama(model_path=self.model_path, n_ctx=690, n_threads=8, n_gpu_layers=32)
+
+        system_prompt = f"Based on the following interaction between a user and an AI assistant, generate a concise topic for the conversation in 2-6 words:\n\nUser: {prompt}\nAssistant: {response}\n\nTopic:"
+
+        topic = llm(
+            system_prompt,
+            max_tokens=10,
+            temperature=0.7,
+            stop=["\\n"],
+            echo=False
+        )
+
+        return topic['choices'][0]['text'].strip()