Update app.py

app.py

@@ -1,19 +1,24 @@
 import json
+from threading import Thread
+
 import gradio as gr
 import spaces
-from huggingface_hub import hf_hub_download
-from llama_cpp import Llama
-from transformers import AutoTokenizer
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    TextIteratorStreamer,
+    FineGrainedFP8Config,
+)
 
 MAX_NEW_TOKENS = 8192
 MODEL_NAME = "Azure99/Blossom-V6.3-36B"
-MODEL_GGUF_REPO = f"{MODEL_NAME}-GGUF"
-MODEL_FILE = "blossom-v6.3-36b-q8_0.gguf"
-MODEL_LOCAL_DIR = "./"
-
-hf_hub_download(repo_id=MODEL_GGUF_REPO, filename=MODEL_FILE, local_dir=MODEL_LOCAL_DIR)
 
-llm: Llama = None
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    torch_dtype="auto",
+    device_map="auto",
+    quantization_config=FineGrainedFP8Config(),
+)
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 
 
@@ -33,27 +38,31 @@ def get_messages(user, history):
 
 @spaces.GPU(duration=120)
 def chat(user, history, temperature, top_p, repetition_penalty):
-    global llm
-    if llm is None:
-        llm = Llama(
-            model_path=MODEL_FILE, n_gpu_layers=-1, flash_attn=True, n_ctx=16384
-        )
+    streamer = TextIteratorStreamer(
+        tokenizer, skip_prompt=True, skip_special_tokens=True
+    )
 
     messages = get_messages(user, history)
     print(f"Messages: {messages}")
-    input_ids = tokenizer.apply_chat_template(messages)
-    generate_config = dict(
+    input_ids = tokenizer.apply_chat_template(messages, return_tensors="pt").to(
+        model.device
+    )
+
+    generation_kwargs = dict(
+        input_ids=input_ids,
+        streamer=streamer,
+        do_sample=True,
+        max_new_tokens=MAX_NEW_TOKENS,
         temperature=temperature,
         top_p=top_p,
-        repeat_penalty=repetition_penalty,
-        top_k=0,
-        stream=True,
-        max_tokens=MAX_NEW_TOKENS,
+        repetition_penalty=repetition_penalty,
     )
 
+    Thread(target=model.generate, kwargs=generation_kwargs).start()
+
     outputs = ""
-    for chunk in llm(input_ids, **generate_config):
-        outputs += chunk["choices"][0]["text"]
+    for new_text in streamer:
+        outputs += new_text
         yield outputs