use ggml model

model.py

@@ -1,32 +1,25 @@
-from threading import Thread
 from typing import Iterator
+from llama_cpp import Llama
+from huggingface_hub import hf_hub_download
 
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
-# Original version
-model_id = "LinkSoul/Chinese-Llama-2-7b"
-# 4 bit version
-# model_id = "LinkSoul/Chinese-Llama-2-7b-4bit"
+def download_model():
+    # See https://github.com/OpenAccess-AI-Collective/ggml-webui/blob/main/tabbed.py
+    # https://huggingface.co/spaces/kat33/llama.cpp/blob/main/app.py
+    print(f"Downloading model: {model_repo}/{model_filename}")
+    file = hf_hub_download(
+            repo_id=model_repo, filename=model_filename
+    )
+    print("Downloaded " + file)
+    return file
 
-if torch.cuda.is_available():
-    if model_id.endswith("4bit"):
-        model = AutoModelForCausalLM.from_pretrained(
-                model_id,
-                load_in_4bit=True,
-                torch_dtype=torch.float16,
-                device_map='auto'
-            )
-    else:
-        model = AutoModelForCausalLM.from_pretrained(
-            model_id,
-            torch_dtype=torch.float16,
-            device_map='auto'
-        )
-else:
-    model = None
-tokenizer = AutoTokenizer.from_pretrained(model_id)
+model_repo = "LinkSoul/Chinese-Llama-2-7b-ggml"
+model_filename = "Chinese-Llama-2-7b.ggmlv3.q4_0.bin"
+# model_filename = "Chinese-Llama-2-7b.ggmlv3.q8_0.bin"
+model_path = download_model()
 
+# load Llama-2
+llm = Llama(model_path=model_path, verbose=False)
 
 
 def get_prompt(message: str, chat_history: list[tuple[str, str]],
@@ -37,11 +30,20 @@ def get_prompt(message: str, chat_history: list[tuple[str, str]],
     texts.append(f'{message.strip()} [/INST]')
     return ''.join(texts)
 
+def generate(prompt, max_new_tokens, temperature, top_p, top_k):
+    return llm(prompt,
+            max_tokens=max_new_tokens,
+            stop=["</s>"],
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            stream=False)
+
 
 def get_input_token_length(message: str, chat_history: list[tuple[str, str]], system_prompt: str) -> int:
     prompt = get_prompt(message, chat_history, system_prompt)
-    input_ids = tokenizer([prompt], return_tensors='np')['input_ids']
-    return input_ids.shape[-1]
+    input_ids = llm.tokenize(prompt.encode('utf-8'))
+    return len(input_ids)
 
 
 def run(message: str,
@@ -52,26 +54,10 @@ def run(message: str,
         top_p: float = 0.95,
         top_k: int = 50) -> Iterator[str]:
     prompt = get_prompt(message, chat_history, system_prompt)
-    inputs = tokenizer([prompt], return_tensors='pt').to('cuda')
-
-    streamer = TextIteratorStreamer(tokenizer,
-                                    timeout=10.,
-                                    skip_prompt=True,
-                                    skip_special_tokens=True)
-    generate_kwargs = dict(
-        inputs,
-        streamer=streamer,
-        max_new_tokens=max_new_tokens,
-        do_sample=True,
-        top_p=top_p,
-        top_k=top_k,
-        temperature=temperature,
-        num_beams=1,
-    )
-    t = Thread(target=model.generate, kwargs=generate_kwargs)
-    t.start()
+    output = generate(prompt, max_new_tokens, temperature, top_p, top_k)
+    yield output['choices'][0]['text']
 
-    outputs = []
-    for text in streamer:
-        outputs.append(text)
-        yield ''.join(outputs)
+    # outputs = []
+    # for resp in streamer:
+    #     outputs.append(resp['choices'][0]['text'])
+    #     yield ''.join(outputs)