Update app.py

app.py

@@ -2,90 +2,98 @@ import os
 import gradio as gr
 from huggingface_hub import hf_hub_download
 from llama_cpp import Llama
-from llama_cpp_agent import LlamaCppAgent, MessagesFormatterType
-from llama_cpp_agent.providers import LlamaCppPythonProvider
-from llama_cpp_agent.chat_history import BasicChatHistory
-from llama_cpp_agent.chat_history.messages import Roles
-
-# Define constants for repository and model details
-REPO_ID = "emizemani/editlyai-gguf"
-MODEL_FILENAME = "editlyai.gguf"
-MODEL_DIR = "./models"
-MODEL_PATH = os.path.join(MODEL_DIR, MODEL_FILENAME)
-
-# Ensure the model directory exists
-os.makedirs(MODEL_DIR, exist_ok=True)
-
-# API token for authenticated access to Hugging Face Hub
-HF_TOKEN = os.getenv("HF_TOKEN")
-
-# Download the model if it's not already present using the HF token for authentication
-if not os.path.exists(MODEL_PATH):
-    hf_hub_download(repo_id=REPO_ID, filename=MODEL_FILENAME, local_dir=MODEL_DIR, use_auth_token=HF_TOKEN)
-
-# Initialize the Llama model for CPU usage
-llm = Llama(
-    model_path=MODEL_PATH,
-    flash_attn=True,  # Adjust this based on your hardware's capability
-    n_gpu_layers=0,  # Using CPU
-    n_batch=1,
-    n_ctx=4096
-)
-
-def respond(message, system_message):
-    # Setup Llama provider and agent
-    provider = LlamaCppPythonProvider(llm)
-    agent = LlamaCppAgent(
-        provider,
-        system_prompt=system_message,
-        predefined_messages_formatter_type=MessagesFormatterType.GEMMA_2,
-        debug_output=True
+
+# 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)[:]
 
-    # Prepare chat history and settings
-    settings = provider.get_provider_default_settings()
-    messages = BasicChatHistory()
-    messages.add_message({'role': Roles.system, 'content': system_message})
-    messages.add_message({'role': Roles.user, 'content': message})
-
-    # Generate response using the Llama model
-    stream = agent.get_chat_response(
-        message,
-        llm_sampling_settings=settings,
-        chat_history=messages,
-        returns_streaming_generator=True,
-        print_output=False
+    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
     )
 
-    # Process the model output
-    output = ""
-    try:
-        for data in stream:
-            if "text" in data:
-                output += data["text"]
-            yield output
-    except Exception as e:
-        yield f"An error occurred: {str(e)}"
-
-# Comprehensive system prompt for Editly AI
-SYSTEM_PROMPT = (
-    "Hello, I am Editly AI, your intelligent text editing assistant. "
-    "I specialize in correcting grammar, enhancing clarity, and refining the style of your text. "
-    "Please type the text you want edited, and I will provide suggestions to improve it."
-)
-
-# Configure Gradio interface
-demo = gr.Interface(
-    fn=respond,
-    inputs=[
-        gr.Textbox(label="Enter your text here", placeholder="Type here...", lines=4),
-        gr.Textbox(value=SYSTEM_PROMPT, label="System message", visible=False),
-    ],
-    outputs="text",
-    title="Editly AI - Text Editing Assistant",
-    description="Interact with Editly AI to refine and enhance your text.",
-    theme="huggingface"
-)
-
-if __name__ == "__main__":
-    demo.launch()
+    demo.launch(show_error=True, share=True)
+
+