Update README.md

README.md

@@ -38,10 +38,12 @@ from transformers import (
     AutoModelForCausalLM,
     BitsAndBytesConfig,
 )
+from peft import LoraConfig, PeftModel
 from transformers import TextStreamer
 
 
-BASE_MODEL = "togepi55/llm-jp-3-13b-it"
+BASE_MODEL = "llm-jp/llm-jp-3-13b"
+PEFT_MODEL = "togepi55/llm-jp-3-13b-it"
 
 tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
 bnb_config = BitsAndBytesConfig(
@@ -51,7 +53,7 @@ bnb_config = BitsAndBytesConfig(
     bnb_4bit_use_double_quant=False,
 )
 
-model = AutoModelForCausalLM.from_pretrained(
+base_model = AutoModelForCausalLM.from_pretrained(
             BASE_MODEL,
             device_map="auto",
             quantization_config=bnb_config,
@@ -59,8 +61,12 @@ model = AutoModelForCausalLM.from_pretrained(
             trust_remote_code=True,
         )
 
+
+model = PeftModel.from_pretrained(base_model, PEFT_MODEL)
+
 streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
 
+
 instruction = "東京は日本の"
 
 prompt = f"<s>以下は、タスクを説明する指示です。要求を適切に満たす応答を書きなさい\n\n### 指示:\n{instruction}\n\n### 応答:\n"