Update app.py

app.py

@@ -2,9 +2,8 @@ import gradio as gr
 import os
 import spaces
 from transformers import AutoTokenizer, TextIteratorStreamer
-from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig
 from threading import Thread
-import torch
+from llama_cpp import Llama
 
 # Set an environment variable
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
@@ -45,15 +44,6 @@ h1 {
 }
 """
 
-# Load the tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained("CyberNative-AI/Colibri_8b_v0.1_gptq_128_4bit")
-model = AutoGPTQForCausalLM.from_quantized("CyberNative-AI/Colibri_8b_v0.1_gptq_128_4bit", dtype=torch.float32, device="cpu")
-
-terminators = [
-    tokenizer.eos_token_id,
-    tokenizer.convert_tokens_to_ids("<|im_end|>")
-]
-
 @spaces.GPU(duration=120)
 def chat_llama3_8b(message: str, 
               history: list, 
@@ -76,24 +66,16 @@ def chat_llama3_8b(message: str,
         conversation.extend([{"role": "user", "content": user}, {"role": "assistant", "content": assistant}])
     conversation.append({"role": "user", "content": message})
 
-    input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt").to(model.device)
-    
-    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
-
-    generate_kwargs = dict(
-        input_ids= input_ids,
-        streamer=streamer,
-        max_new_tokens=max_new_tokens,
-        do_sample=True,
-        top_p=0.7,
-        temperature=temperature,
-        eos_token_id=terminators,
+    llm = Llama.from_pretrained(
+        repo_id="CyberNative-AI/Colibri_8b_v0.1_q5_gguf",
+        filename="*Q5_K_M.gguf",
+        chat_format="chatml",
+        verbose=False,
+        max_tokens=max_new_tokens,
+        stop=["<|im_end|>"]
     )
-    # This will enforce greedy generation (do_sample=False) when the temperature is passed 0, avoiding the crash.             
-    if temperature == 0:
-        generate_kwargs['do_sample'] = False
-        
-    t = Thread(target=model.generate, kwargs=generate_kwargs)
+    
+    t = Thread(target=llm.create_chat_completion, messages=conversation, temperature=temperature)
     t.start()
 
     outputs = []