Update README.md

README.md

@@ -138,32 +138,35 @@ deepspeed --num_gpus 8 src/train_bash.py \
 # 2. Usage
 
 ```python
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from llama_cpp import Llama
 
-model_id = "shenzhi-wang/Llama3-8B-Chinese-Chat"
-
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(
-    model_id, torch_dtype="auto", device_map="auto"
+model = Llama(
+    "/Your/Path/To/GGUF/File",
+    verbose=False,
+    n_gpu_layers=-1,
 )
 
+system_prompt = "You are a helpful assistant."
+
+def generate_reponse(_model, _messages, _max_tokens=8192):
+    _output = _model.create_chat_completion(
+        _messages,
+        stop=["<|eot_id|>", "<|end_of_text|>"],
+        max_tokens=_max_tokens,
+    )["choices"][0]["message"]["content"]
+    return _output
+
+# The following are some examples
+
 messages = [
+    {
+        "role": "system",
+        "content": system_prompt,
+    },
     {"role": "user", "content": "写一首诗吧"},
 ]
 
-input_ids = tokenizer.apply_chat_template(
-    messages, add_generation_prompt=True, return_tensors="pt"
-).to(model.device)
-
-outputs = model.generate(
-    input_ids,
-    max_new_tokens=8192,
-    do_sample=True,
-    temperature=0.6,
-    top_p=0.9,
-)
-response = outputs[0][input_ids.shape[-1]:]
-print(tokenizer.decode(response, skip_special_tokens=True))
+print(generate_reponse(model, messages))
 ```
 
 # 3. Examples