|
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): |
|
|
|
pass |
