Update README.md

README.md

@@ -39,44 +39,37 @@ llama3.2-3B가 나왔는데 한국어가 포함 안되었다구?? 이번 Bllosso
 
 
 ```python
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from llama_cpp import Llama
+from transformers import AutoTokenizer
 
 model_id = 'Bllossom/llama-3.2-Korean-Bllossom-3B'
-
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(
-    model_id,
-    torch_dtype=torch.bfloat16,
-    device_map="auto",
+model = Llama(
+    model_path='llama-3.2-Korean-Bllossom-3B-gguf-Q4_K_M.gguf'
 )
+
 instruction = "철수가 20개의 연필을 가지고 있었는데 영희가 절반을 가져가고 민수가 남은 5개를 가져갔으면 철수에게 남은 연필의 갯수는 몇개인가요?"
 
 messages = [
     {"role": "user", "content": f"{instruction}"}
     ]
 
-input_ids = tokenizer.apply_chat_template(
-    messages,
-    add_generation_prompt=True,
-    return_tensors="pt"
-).to(model.device)
-
-terminators = [
-    tokenizer.convert_tokens_to_ids("<|end_of_text|>"),
-    tokenizer.convert_tokens_to_ids("<|eot_id|>")
-]
-
-outputs = model.generate(
-    input_ids,
-    max_new_tokens=1024,
-    eos_token_id=terminators,
-    do_sample=True,
-    temperature=0.6,
-    top_p=0.9
+prompt = tokenizer.apply_chat_template(
+    messages, 
+    tokenize = False,
+    add_generation_prompt=True
 )
 
-print(tokenizer.decode(outputs[0][input_ids.shape[-1]:], skip_special_tokens=True))
+generation_kwargs = {
+    "max_tokens":512,
+    "stop":["<|eot_id|>"],
+    "echo":True,
+    "top_p":0.9,
+    "temperature":0.6,
+}
+
+resonse_msg = model(prompt, **generation_kwargs)
+print(resonse_msg['choices'][0]['text'][len(prompt):])
 ```
 ```
 철수가 20개의 연필을 가지고 있었고 영희가 절반을 가져가면, 영희가 가져간 연필의 갯수는 20 / 2 = 10개입니다.