84 lines
2.6 KiB
Python
84 lines
2.6 KiB
Python
import sys
|
|
from core import llm_room as ap
|
|
from config.llm import (
|
|
MODEL, N_THREADS, N_CTX, N_GPU_LAYERS, MAX_TOKENS, STREAM,
|
|
TEMPERATURE, TOP_P, TOP_K, SYSTEM_PROMPT, CHAT_TEMPLATE,
|
|
)
|
|
from llama_cpp import Llama
|
|
|
|
def load_llm():
|
|
print("Loading model...")
|
|
llm = Llama(
|
|
model_path = str(MODEL),
|
|
n_ctx = N_CTX,
|
|
n_threads = N_THREADS,
|
|
n_gpu_layers = N_GPU_LAYERS,
|
|
chat_format = CHAT_TEMPLATE or None,
|
|
verbose = False,
|
|
)
|
|
print("Model ready!")
|
|
return llm
|
|
|
|
def chat_once(llm, messages, max_tokens):
|
|
if STREAM:
|
|
text = ""
|
|
for chunk in llm.create_chat_completion(
|
|
messages,
|
|
stream = True,
|
|
temperature = TEMPERATURE,
|
|
top_p = TOP_P,
|
|
top_k = TOP_K,
|
|
max_tokens = max_tokens,
|
|
):
|
|
delta = chunk["choices"][0]["delta"].get("content")
|
|
if delta:
|
|
print(delta, end="", flush=True)
|
|
text += delta
|
|
print()
|
|
return text
|
|
|
|
resp = llm.create_chat_completion(
|
|
messages,
|
|
stream = False,
|
|
temperature = TEMPERATURE,
|
|
top_p = TOP_P,
|
|
top_k = TOP_K,
|
|
max_tokens = max_tokens,
|
|
)
|
|
return resp["choices"][0]["message"]["content"]
|
|
|
|
def llm_run(args):
|
|
|
|
if not MODEL.is_file():
|
|
print(f"Model not found: {MODEL}. Download it first (see README).", file=sys.stderr)
|
|
return
|
|
|
|
llm = load_llm()
|
|
system = args.system if args.system is not None else SYSTEM_PROMPT
|
|
max_tokens = args.max_tokens if args.max_tokens is not None else MAX_TOKENS
|
|
|
|
if args.prompt is not None:
|
|
messages = [{"role": "system", "content": system}] if system else []
|
|
messages.append({"role": "user", "content": args.prompt})
|
|
chat_once(llm, messages, max_tokens)
|
|
return
|
|
|
|
messages = [{"role": "system", "content": system}] if system else []
|
|
print("Interactive chat. Type 'exit' or 'quit' to leave.")
|
|
try:
|
|
while True:
|
|
user = input("\nYou : ").strip()
|
|
if user.lower() in ("exit", "quit", "q"):
|
|
print("Bye.")
|
|
break
|
|
if not user:
|
|
continue
|
|
messages.append({"role": "user", "content": user})
|
|
print("Assistant: ", end="", flush=True)
|
|
reply = chat_once(llm, messages, max_tokens)
|
|
messages.append({"role": "assistant", "content": reply})
|
|
except (KeyboardInterrupt, EOFError):
|
|
print("\nBye.")
|
|
|
|
if __name__ == "__main__":
|
|
llm_run( ap.parser.parse_args() ) |