xsa-dev's picture
init llama 6
0d9e7b2
import gradio as gr
import copy
from llama_cpp import Llama
from huggingface_hub import hf_hub_download # load from huggingfaces
CONST_REPO_ID = "TheBloke/Llama-2-7B-Chat-GGML"
CONST_FILENAME = "llama-2-7b-chat.ggmlv3.q6_K.bin"
N_CTX = 4096
llm = Llama(model_path=hf_hub_download(
repo_id=CONST_REPO_ID,
filename=CONST_FILENAME),
n_ctx=N_CTX
)
history = N_CTX
pre_prompt = \
" The user and the AI are having a conversation : <|endoftext|> \n"
def generate_text(input_text, history):
temp = ""
if history == []:
input_text_with_history = f"SYSTEM:{pre_prompt}" + \
"\n" + f"USER: {input_text} " + "\n" + " ASSISTANT:"
else:
input_text_with_history = f"{history[-1][1]}" + "\n"
input_text_with_history += f"USER: {input_text}" + "\n" + " ASSISTANT:"
output = llm(input_text_with_history, max_tokens=4096, stop=[
"<|prompter|>", "<|endoftext|>", "<|endoftext|> \n",
"ASSISTANT:", "USER:", "SYSTEM:"], stream=True
)
for out in output:
stream = copy.deepcopy(out)
temp += stream["choices"][0]["text"]
yield temp
history = ["init", input_text_with_history]
demo = gr.ChatInterface(generate_text,
title=f"Lama2 on CPU: {CONST_FILENAME}",
description=f"Running Llama2 with llama_cpp: \
\r\n<i>{CONST_REPO_ID} {CONST_FILENAME}</i>",
examples=["Hi!",
"Does it hard to be machine?",
"When i am need a doctor?",
"Ты говоришь по русски? Я злой."
],
cache_examples=True,
undo_btn="Undo",
clear_btn="Clear")
demo.queue(concurrency_count=10, max_size=50)
demo.launch()