Update app.py

app.py

@@ -1,18 +1,21 @@
 import spaces
 import gradio as gr
-from llama_cpp import Llama
-from llama_cpp_agent import LlamaCppAgent
-from llama_cpp_agent import MessagesFormatterType
-from llama_cpp_agent.providers import LlamaCppPythonProvider
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from threading import Thread
 
 """
 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=4, n_gpu_layers=33, n_ctx=8192, verbose=False)
-
-provider = LlamaCppPythonProvider(llama_model)
 
+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
 
 @spaces.GPU
 def respond(
@@ -23,21 +26,34 @@ def respond(
     temperature,
     top_p,
 ):
-    agent = LlamaCppAgent(
-        provider,
-        system_prompt=system_message,
-        predefined_messages_formatter_type=MessagesFormatterType.LLAMA_3,
-        debug_output=True
-    )
+    torch.set_default_device("cuda")
 
-    settings = provider.get_provider_default_settings()
+    model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.2", torch_dtype="auto",load_in_4bit=True,trust_remote_code=True)
+    tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.2", trust_remote_code=True)
+ 
+    history_transformer_format = history + [[message, ""]]
 
-    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()
+    messages = system_prompt + "".join(["".join(["\n[INST]" + item[0], "[/INST]\n" + item[1] + "</s>"]) 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
 
 """
 For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
@@ -45,7 +61,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 helpful assistant.", label="System message"),
+        gr.Textbox(value="You are a friendly Chatbot.", 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(