Update app.py

app.py

@@ -1,21 +1,14 @@
 import spaces
 import gradio as gr
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-from threading import Thread
+from huggingface_hub import hf_hub_download
+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")
+subprocess.run('pip install llama-cpp-python', env={'CMAKE_ARGS': "-DLLAMA_CUBLAS=ON", 'FORCE_CMAKE': '1'}, shell=True)
 
-def format_prompt(message, history):
-  prompt = "<s>"
-  for user_prompt, bot_response in history:
-    prompt += f"[INST] {user_prompt} [/INST]"
-    prompt += f" {bot_response}</s> "
-  prompt += f"[INST] {message} [/INST]"
-  return prompt
+hf_hub_download(repo_id="TheBloke/Mistral-7B-Instruct-v0.2-GGUF", filename="mistral-7b-instruct-v0.2.Q6_K.gguf")
 
 @spaces.GPU
 def respond(
@@ -26,50 +19,30 @@ def respond(
     temperature,
     top_p,
 ):
-    torch.set_default_device("cuda")
+    llama_model = Llama(r"mistral-7b-instruct-v0.2.Q6_K.gguf", n_batch=1024, n_threads=0,            n_gpu_layers=33, n_ctx=8192, verbose=False)
 
-    tokenizer = AutoTokenizer.from_pretrained(
-        "Weyaxi/Einstein-v6.1-Llama3-8B",
-        trust_remote_code=True
-    )
-    model = AutoModelForCausalLM.from_pretrained(
-        "Weyaxi/Einstein-v6.1-Llama3-8B",
-        torch_dtype="auto",
-        load_in_4bit=True,
-        trust_remote_code=True
-    )
-    history_transformer_format = history + [[message, ""]]
+    provider = LlamaCppPythonProvider(llama_model)
 
-    system_prompt = "<|im_start|>system\nYou are Einstein, a helpful AI assistant.<|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
+    agent = LlamaCppAgent(
+      provider,
+      system_prompt=f"{system_message}",
+      predefined_messages_formatter_type=MessagesFormatterType.MISTRAL,
+      debug_output=True
     )
-    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
 
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
+    settings = provider.get_provider_default_settings()
+    settings.stream = True
+    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()
+
 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(
@@ -80,11 +53,7 @@ demo = gr.ChatInterface(
             label="Top-p (nucleus sampling)",
         ),
     ],
-    theme=gr.themes.Soft(primary_hue="green", secondary_hue="indigo", neutral_hue="zinc",font=[gr.themes.GoogleFont("Exo 2"), "ui-sans-serif", "system-ui", "sans-serif"]).set(
-        block_background_fill_dark="*neutral_800"
-    )
 )
 
-
 if __name__ == "__main__":
     demo.launch()