Update app.py

app.py

@@ -1,13 +1,18 @@
 import spaces
 import gradio as gr
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-from threading import Thread
+from llama_cpp import Llama
+from llama_cpp_agent import LlamaCppAgent
+from llama_cpp_agent import MessagesFormatterType
+from llama_cpp_agent.providers import LlamaCppPythonProvider
 
 """
 For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
 """
 # client = InferenceClient("cognitivecomputations/dolphin-2.8-mistral-7b-v02")
+llama_model = Llama(r"Meta-Llama-3-8B.Q5_K_M.gguf", n_batch=1024, n_threads=10, n_gpu_layers=33, n_ctx=8192, verbose=False)
+
+provider = LlamaCppPythonProvider(llama_model)
+
 
 @spaces.GPU
 def respond(
@@ -18,42 +23,21 @@ def respond(
     temperature,
     top_p,
 ):
-    torch.set_default_device("cuda")
-
-    tokenizer = AutoTokenizer.from_pretrained(
-        "cognitivecomputations/dolphin-2.8-mistral-7b-v02",
-        trust_remote_code=True
-    )
-    model = AutoModelForCausalLM.from_pretrained(
-        "cognitivecomputations/dolphin-2.8-mistral-7b-v02",
-        torch_dtype="auto",
-        load_in_4bit=True,
-        trust_remote_code=True
+    agent = LlamaCppAgent(
+        provider,
+        system_prompt=system_message,
+        predefined_messages_formatter_type=MessagesFormatterType.MISTRAL,
+        debug_output=True
     )
-    history_transformer_format = history + [[message, ""]]
 
-    system_prompt = f"<|im_start|>system\n{system_message}.<|im_end|>"
-    messages = system_prompt + "".join(["".join(["\n<|im_start|>user\n" + item[0], "<|im_end|>\n<|im_start|>assistant\n" + item[1]]) for item in history_transformer_format])
-    input_ids = tokenizer([messages], return_tensors="pt").to('cuda')
-    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-    generate_kwargs = dict(
-        input_ids,
-        streamer=streamer,
-        max_new_tokens=max_tokens,
-        do_sample=True,
-        top_p=top_p,
-        top_k=50,
-        temperature=temperature,
-        num_beams=1
-    )
-    t = Thread(target=model.generate, kwargs=generate_kwargs)
-    t.start()
-    partial_message = ""
-    for new_token in streamer:
-        partial_message += new_token
-        if '<|im_end|>' in partial_message:
-            break
-        yield partial_message
+    settings = provider.get_provider_default_settings()
+
+    settings.max_tokens = max_tokens
+    settings.temperature = temperature
+    settings.top_p = top_p
+
+    agent_output = agent.get_chat_response(message, llm_sampling_settings=settings)
+    yield agent_output.strip()
 
 """
 For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
@@ -61,7 +45,7 @@ For information on how to customize the ChatInterface, peruse the gradio docs: h
 demo = gr.ChatInterface(
     respond,
     additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
+        gr.Textbox(value="You are a helpful assistant.", label="System message"),
         gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
         gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
         gr.Slider(