Update README.md

README.md

@@ -106,9 +106,10 @@ Dialogue History:
 **Example Code**
 
 ```python
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 from prompts import overall_evaluation_gen_prompt ## same prompt on "Instruction for TD-Llama-General"
-import json,re
+import json,re,torch
+from peft import PeftModel, PeftConfig
 
 def normalize_action(text):
     lower_text = text.lower()
@@ -182,8 +183,18 @@ def overall_form(dialogue_idx,full_dialogue,tokenizer_inf,is_val=False):
     return label
 
 device = "cuda:0"
-model = AutoModelForCausalLM.from_pretrained("HOLILAB/td-llama-general",device_map={"": device})
+base_model = AutoModelForCausalLM.from_pretrained(
+    "meta-llama/Llama-3.1-8B-Instruct",
+    quantization_config=BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=torch.bfloat16,
+        bnb_4bit_use_double_quant=True,
+    ),
+    device_map={"": device},
+)
 tokenizer = AutoTokenizer.from_pretrained("HOLILAB/td-llama-general")
+model = PeftModel.from_pretrained(base_model, "HOLILAB/td-llama-general")
 
 with open("test_dialogue_overall_obs.json",'r') as f: # https://drive.google.com/drive/folders/17IASNvCcJRkHlg2tzMPRlpPbo_MlHs94?hl=ko
     dial_list = json.load(f)