Update README.md

README.md

@@ -26,34 +26,8 @@ base_model: meta-llama/llama-2-7b-hf
 ## Instruction format
 
 ```python
+from transformers import  AutoModelForCausalLM, LlamaTokenizer, BitsAndBytesConfig, TextStreamer, StoppingCriteria, StoppingCriteriaList
 import torch
-from transformers import (
-    AutoModelForCausalLM,
-    AutoTokenizer,
-    TextStreamer,
-    StoppingCriteria,
-    StoppingCriteriaList,
-    BitsAndBytesConfig,
-)
-
-device = "cuda"  # the device to load the model onto
-model_name = "willnguyen/lacda-2-7B-chat-v0.1"
-
-bnb_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_compute_dtype=torch.float16,
-    bnb_4bit_use_double_quant=True,
-)
-model = AutoModelForCausalLM.from_pretrained(
-    model_name, load_in_4bit=True, torch_dtype=torch.float16, quantization_config=bnb_config, device_map="auto"
-)
-
-tokenizer = AutoTokenizer.from_pretrained(
-    model_name, cache_dir=None, use_fast=False, padding_side="right", tokenizer_type="llama"
-)
-tokenizer.pad_token_id = 0
-
 
 class StopTokenCriteria(StoppingCriteria):
     def __init__(self, stop_tokens, tokenizer, prompt_length):
@@ -69,20 +43,33 @@ class StopTokenCriteria(StoppingCriteria):
 
     def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
         is_done = False
-        tokens = tokenizer.decode(input_ids[0])[self.prompt_length :]
+        tokens = tokenizer.decode(input_ids[0])[self.prompt_length:]
         for st in self.stop_tokens:
             if st in tokens:
                 is_done = True
                 break
         return is_done
 
+model_name = "willnguyen/lacda-2-7B-chat-v0.1"
+tokenizer = LlamaTokenizer.from_pretrained(
+        model_name,
+        use_fast=False,
+        padding_side="right",
+        tokenizer_type='llama',
+)
+tokenizer.pad_token_id = 0
 
-# Pass in a prompt and infer with the model
-prompt = "<s> [INST] Who is Ho Chi Minh? [/INST]"
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+     device_map="auto",
+    torch_dtype=torch.float16,
+)
+
+prompt = "<s> [INST] who is Hồ Chí Minh [/INST]"
 
 stopping_criteria = StoppingCriteriaList([StopTokenCriteria(["[INST]", "[/INST]"], tokenizer, len(prompt))])
 with torch.inference_mode():
-    input_ids = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).input_ids.to(device)
+    input_ids = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).input_ids.to('cuda')
     streamer = TextStreamer(tokenizer)
     _ = model.generate(
         input_ids=input_ids,
@@ -94,7 +81,6 @@ with torch.inference_mode():
         repetition_penalty=1.0,
         use_cache=True,
         streamer=streamer,
-        stopping_criteria=stopping_criteria,
+        stopping_criteria=stopping_criteria
     )
-
 ```