File size: 7,028 Bytes
8e9f89e 8dc84a2 57fda94 8e9f89e 71e3aec 8dc84a2 57fda94 a716951 57fda94 8e9f89e 8dc84a2 8e9f89e 8dc84a2 8e9f89e 8dc84a2 8e9f89e 8dc84a2 8e9f89e 9370091 8dc84a2 9370091 8dc84a2 9370091 4ec5dc5 9370091 4ec5dc5 9370091 4ec5dc5 9370091 8e9f89e 8dc84a2 9bb723f 8e9f89e 8dc84a2 8e9f89e 2f6a972 8e9f89e 9370091 bcef7b8 a716951 9370091 a716951 57fda94 bcef7b8 57fda94 a716951 8e9f89e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 |
import aiosqlite
from datetime import datetime
from ctransformers import AutoModelForCausalLM
import asyncio
class PMBL:
def __init__(self, model_path, gpu_layers=50):
self.model_path = model_path
self.gpu_layers = gpu_layers
self.db_name = 'chat_history.db'
self.init_db_lock = asyncio.Lock()
async def init_db(self):
async with self.init_db_lock:
async with aiosqlite.connect(self.db_name) as db:
await db.execute('''CREATE TABLE IF NOT EXISTS chats
(id INTEGER PRIMARY KEY AUTOINCREMENT,
timestamp TEXT,
prompt TEXT,
response TEXT,
topic TEXT)''')
await db.commit()
async def get_chat_history(self, mode="full", user_message=""):
await self.init_db()
async with aiosqlite.connect(self.db_name) as db:
if mode == "full":
async with db.execute("SELECT prompt, response FROM chats ORDER BY id") as cursor:
rows = await cursor.fetchall()
history = []
for row in rows:
history.append({"role": "user", "content": row[0]})
history.append({"role": "PMB", "content": row[1]})
else:
async with db.execute("SELECT id, prompt, response FROM chats WHERE topic != 'Untitled'") as cursor:
chats = await cursor.fetchall()
relevant_chat_id = await self.find_relevant_chat(chats, user_message)
if relevant_chat_id:
async with db.execute("SELECT prompt, response FROM chats WHERE id = ?", (relevant_chat_id,)) as cursor:
row = await cursor.fetchone()
history = [
{"role": "user", "content": row[0]},
{"role": "PMB", "content": row[1]}
]
else:
history = []
return history
async def find_relevant_chat(self, chats, user_message):
max_score = 0
relevant_chat_id = None
for chat in chats:
chat_id, prompt, response = chat
score = self.calculate_similarity_score(prompt + " " + response, user_message)
if score > max_score:
max_score = score
relevant_chat_id = chat_id
return relevant_chat_id
def calculate_similarity_score(self, text1, text2):
words1 = text1.lower().split()
words2 = text2.lower().split()
score = 0
for i in range(len(words1) - 1):
if words1[i] in words2 and words1[i + 1] in words2:
score += 1
return score
async def save_chat_history(self, prompt, response):
await self.init_db()
async with aiosqlite.connect(self.db_name) as db:
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
await db.execute("INSERT INTO chats (timestamp, prompt, response, topic) VALUES (?, ?, ?, 'Untitled')", (timestamp, prompt, response))
await db.commit()
async def generate_response(self, prompt, history, mode):
history.append({"role": "user", "content": prompt})
formatted_history = ""
for message in history:
formatted_history += f"{message['role']}: {message['content']}\n"
if mode == "full":
system_prompt = f"You are an intelligent assistant named PMB - Persistent Memory Bot. You answer any 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.\n\n{formatted_history}\nPMB:"
else:
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:"
n_ctx = self.calculate_context(system_prompt, formatted_history)
llm = AutoModelForCausalLM.from_pretrained(
self.model_path,
model_type="llama",
gpu_layers=self.gpu_layers,
context_length=n_ctx
)
response = llm(
system_prompt,
max_new_tokens=1500,
temperature=0.7,
stop=["</s>", "\nUser:", "\nuser:", "\nSystem:", "\nsystem:"],
stream=True
)
response_text = ""
for chunk in response:
response_text += chunk
yield chunk
await self.save_chat_history(prompt, response_text)
def calculate_context(self, system_prompt, formatted_history):
system_prompt_tokens = len(system_prompt) // 4
history_tokens = len(formatted_history) // 4
max_response_tokens = 1500
context_ceiling = 32690
available_tokens = context_ceiling - system_prompt_tokens - max_response_tokens
if history_tokens <= available_tokens:
return system_prompt_tokens + history_tokens + max_response_tokens
else:
return context_ceiling
async def sleep_mode(self):
await self.init_db()
async with aiosqlite.connect(self.db_name) as db:
async with db.execute("SELECT id, prompt, response FROM chats WHERE topic = 'Untitled'") as cursor:
untitled_chats = await cursor.fetchall()
for chat in untitled_chats:
chat_id, prompt, response = chat
topic = await self.generate_topic(prompt, response)
await db.execute("UPDATE chats SET topic = ? WHERE id = ?", (topic, chat_id))
await db.commit()
async def generate_topic(self, prompt, response):
llm = AutoModelForCausalLM.from_pretrained(
self.model_path,
model_type="llama",
gpu_layers=self.gpu_layers,
context_length=2960
)
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:\n\nUser: {prompt}\nAssistant: {response}\n\nTopic:"
topic = llm(
system_prompt,
max_new_tokens=12,
temperature=0,
stop=["\n"]
)
return topic.strip()
async def close(self):
# Implement any cleanup operations here if needed
pass |