File size: 3,169 Bytes
26f89b1 f57f476 26f89b1 c52aaee 26f89b1 c52aaee 26f89b1 c52aaee 26f89b1 c52aaee | 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 | import os
import gradio as gr
from huggingface_hub import hf_hub_download
from llama_cpp import Llama
# Define the system prompt for Editly AI
SYSTEM_PROMPT = "You are Editly AI, an English-speaking assistant specialized in editing and improving text. Communicate with users to enhance their writing. You will answer queries about editing only."
def get_message_tokens(model, role, content):
content = f"{role}\n{content}\n</s>"
content = content.encode("utf-8")
return model.tokenize(content, special=True)
def get_system_tokens(model):
system_message = {"role": "system", "content": SYSTEM_PROMPT}
return get_message_tokens(model, **system_message)
def load_model(
directory: str = ".",
model_name: str = "editlyai.gguf",
repo_id: str = "emizemani/editlyai-gguf"
):
final_model_path = os.path.join(directory, model_name)
print("Downloading model...")
if not os.path.exists(final_model_path):
hf_hub_download(
repo_id=repo_id,
filename=model_name,
local_dir=directory,
use_auth_token=os.getenv('HF_TOKEN') # Ensure you have the HF_TOKEN set up in your environment
)
print("Model downloaded and ready to use.")
model = Llama(
model_path=final_model_path,
n_ctx=1024
)
print("Model loaded!")
return model
# Initialize the model
MODEL = load_model()
def user(message, history):
new_history = history + [[message, None]]
return "", new_history
def bot(history, system_prompt):
model = MODEL
tokens = get_system_tokens(model)[:]
for user_message, bot_message in history[:-1]:
message_tokens = get_message_tokens(model=model, role="user", content=user_message)
tokens.extend(message_tokens)
if bot_message:
message_tokens = get_message_tokens(model=model, role="bot", content=bot_message)
tokens.extend(message_tokens)
last_user_message = history[-1][0]
message_tokens = get_message_tokens(model=model, role="user", content=last_user_message)
tokens.extend(message_tokens)
role_tokens = model.tokenize("bot\n".encode("utf-8"), special=True)
tokens.extend(role_tokens)
generator = model.generate(tokens)
partial_text = ""
for token in generator:
if token == model.token_eos():
break
partial_text += model.detokenize([token]).decode("utf-8", "ignore")
history[-1][1] = partial_text
yield history
# Gradio interface setup
with gr.Blocks(theme=gr.themes.Soft()) as demo:
gr.Markdown(f"<h1>Editly AI - Text Editing Assistant</h1>")
with gr.Row():
system_prompt_box = gr.Textbox(value=SYSTEM_PROMPT, label="System Prompt", interactive=False)
chatbot = gr.Chatbot(label="Conversation")
with gr.Row():
msg = gr.Textbox(label="Send a message")
submit = gr.Button("Send")
submit.click(
fn=user,
inputs=[msg, chatbot],
outputs=[msg, chatbot]
).success(
fn=bot,
inputs=[chatbot, system_prompt_box],
outputs=chatbot
)
demo.launch(show_error=True, share=True)
|