Update app.py

app.py

@@ -1,89 +1,37 @@
 import gradio as gr
 from llama_cpp import Llama
+from huggingface_hub import hf_hub_download
 import os
-import requests
 
-MODEL_PATH = "irixium-12b-model_stock-q4_k_m.gguf"
-DEFAULT_SYSTEM_PROMPT = "You are a helpful assistant."
-
-def download_model(url, save_path):
-    try:
-        response = requests.get(url, stream=True)
-        response.raise_for_status()
-        with open(save_path, "wb") as file:
-            for chunk in response.iter_content(chunk_size=8192):
-                file.write(chunk)
-        return "Model downloaded successfully."
-    except Exception as e:
-        return f"Error: {e}"
-
-def load_model(model_path):
-    try:
-        llm = Llama(model_path, n_threads=2, n_gpu_layers=0) #force cpu, and set threads.
-        return llm
-    except Exception as e:
-        return f"Error: {e}"
-
-def apply_chat_template(model_name, messages, system_prompt):
-    model_name_lower = model_name.lower()
-    if "llama-2" in model_name_lower:
-        template = "<s>[INST] <<SYS>>\n{system_prompt}\n<</SYS>>\n\n{user_message} [/INST] {assistant_message}</s>"
-    elif "mistral" in model_name_lower:
-        template = "<s>[INST] {user_message} [/INST] {assistant_message}</s>"
-    else:
-        template = "{user_message} {assistant_message}"
-
-    formatted_messages = []
-    for message in messages:
-        if message["role"] == "system":
-            system_prompt = message["content"]
-        elif message["role"] == "user":
-            formatted_messages.append(template.format(system_prompt=system_prompt, user_message=message["content"], assistant_message=""))
-        elif message["role"] == "assistant":
-            if formatted_messages:
-                formatted_messages[-1] += message["content"]
-            else:
-                formatted_messages.append(message["content"])
-    return "".join(formatted_messages)
-
-def generate_response(prompt, model):
-    if isinstance(model, str):
-        return model
-    try:
-        output = model(prompt, max_tokens=256)
-        return output["choices"][0]["text"].strip()
-    except Exception as e:
-        return f"Error: {e}"
-
-def inference(message, history, model_url, system_prompt):
-    if model_url and not os.path.exists(MODEL_PATH):
-        download_result = download_model(model_url, MODEL_PATH)
-        if "Error" in download_result:
-            return history + [{"role": "assistant", "content": download_result}], history
-
-    llm = load_model(MODEL_PATH)
-    if isinstance(llm, str):
-        return history + [{"role": "assistant", "content": llm}], history
-
-    messages = [{"role": "system", "content": system_prompt}]
-    for item in history:
-        messages.append(item)
-    messages.append({"role": "user", "content": message})
-
-    prompt = apply_chat_template(llm.model_path, messages, system_prompt)
-    response = generate_response(prompt, llm)
-    history.append({"role": "assistant", "content": response})
-    return history, history
-
-with gr.Blocks() as iface:
-    model_url_input = gr.Textbox(label="Model URL (GGUF)", placeholder="Enter GGUF model URL...")
-    system_prompt_input = gr.Textbox(label="System Prompt", value=DEFAULT_SYSTEM_PROMPT, lines=3)
-    chatbot = gr.Chatbot(type="messages")
-    message = gr.Textbox(label="Message")
-    send_button = gr.Button("Send")
-    state = gr.State([])
-
-    send_button.click(inference, inputs=[message, state, model_url_input, system_prompt_input], outputs=[chatbot, state])
-    message.submit(inference, inputs=[message, state, model_url_input, system_prompt_input], outputs=[chatbot, state])
+MODEL_REPO = "DreadPoor/Irixium-12B-Model_Stock-Q4_K_M-GGUF"
+MODEL_FILENAME = "irixium-12b-model_stock-q4_k_m.gguf"
+MODEL_PATH = "./" + MODEL_FILENAME
+
+if not os.path.exists(MODEL_PATH):
+    hf_hub_download(
+        repo_id=MODEL_REPO,
+        filename=MODEL_FILENAME,
+        repo_type="model",
+        local_dir=".",
+    )
+
+llm = Llama(
+    model_path=MODEL_PATH,
+    n_ctx=4096,
+    n_threads=2,
+    n_threads_batch=2,
+    verbose=False,
+)
+
+def generate_response(message, history):
+    prompt = f"{message}"
+    output = llm(prompt, max_tokens=128, echo=False)
+    return output["choices"][0]["text"].strip()
+
+iface = gr.ChatInterface(
+    fn=generate_response,
+    title="llama.cpp Chat",
+    description="Chat with a GGUF model.",
+)
 
 iface.launch()