demo-llm / chat.py
Petro
model
036f518
from llama_cpp import Llama
from llama_cpp import ChatCompletionRequestMessage as Message
from llama_cpp import ChatCompletionRequestSystemMessage as SystemMessage
from llama_cpp import ChatCompletionRequestAssistantMessage as AssistantMessage
from llama_cpp import ChatCompletionRequestUserMessage as UserMessage
SYSTEM = 'system'
USER = 'user'
ASSISTANT = 'assistant'
EXIT = 'exit'
model_path = "zephyr-7b-beta.Q4_K_S.gguf"
llm = Llama(model_path=model_path, n_ctx=512, max_answer_len=100) # Set chat_format according to the model you are using
class Chat:
def __init__(self, model: Llama) -> None:
self.model: Llama = model
self.messages: list[Message] = [
SystemMessage(
role=SYSTEM,
content='You are a helpful developer assistant, answer all the questions correctly and concisely.'
),
AssistantMessage(role=ASSISTANT, content='Hello, do you have any question?'),
]
def send_message(self, content: str):
new_message = UserMessage(role=USER, content=content)
self.messages.append(new_message)
def generate_reply(self) -> str:
response = self.model.create_chat_completion(
messages=self.messages,
temperature=0.7,
top_p=0.9,
top_k=20,
max_tokens=128
)
response_content = response['choices'][0]['message']
self.messages.append(AssistantMessage(role=ASSISTANT, content=response_content))
return response_content