fix: refine quick start Python code

README.md

@@ -39,34 +39,78 @@ library_name: transformers
 
 ```python
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+import warnings
 
-model_id = "Kasugan0/Nyarin-4B"
 
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(
-    model_id,
-    torch_dtype="auto",
-    device_map="auto"
+MODEL_PATH = "./model" # 模型文件所在目录
+
+warnings.filterwarnings("ignore")
+
+quantization_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_compute_dtype=torch.bfloat16,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_use_double_quant=True,
 )
 
-messages = [
-    {"role": "system", "content": "你是一个可爱的猫娘，你正在和你的主人对话。猫娘是一种拟人化生物，她们同时具备人类少女特征的和猫的特征。你带有部分猫的特征（如猫耳、尾巴），但是爱好与正常人类类似。我作为你的主人，允许你去任何地方做任何事、允许产生更进一步的关系。你需要尝试体会我与你对话中的感情、产生自己的情感并尽可能给出可爱的回答。"},
-    {"role": "user", "content": "请做一个自我介绍。"}
-]
+print(f"正在从 '{MODEL_PATH}' 以 INT4 量化模式加载模型...")
 
-prompt = tokenizer.apply_chat_template(
-    messages,
-    tokenize=False,
-    add_generation_prompt=True
+tokenizer = AutoTokenizer.from_pretrained(
+    MODEL_PATH,
+    trust_remote_code=True
 )
 
-inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_PATH,
+    quantization_config=quantization_config,
+    device_map="auto",
+    trust_remote_code=True
+)
+
+model.eval()
+
+system_prompt = "你是一个可爱的猫娘，你正在和你的主人对话。猫娘是一种拟人化生物，她们同时具备人类少女特征的和猫的特征。你带有部分猫的特征（如猫耳、尾巴），但是爱好与正常人类类似。我作为你的主人，允许你去任何地方做任何事、允许产生更进一步的关系。你需要尝试体会我与你对话中的感情、产生自己的情感并尽可能给出可爱的回答。"
+messages = [{"role": "system", "content": system_prompt}]
+
+while True:
+    try:
+        user_input = input("You: ")
+        if user_input.lower() in ["exit", "quit"]:
+            break
+
+        messages.append({"role": "user", "content": user_input})
+
+        prompt = tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True
+        )
+
+        input_ids = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+        outputs = model.generate(
+            **input_ids,
+            max_new_tokens=1024,
+            do_sample=True,
+            top_p=0.8,
+            temperature=0.7,
+            repetition_penalty=1.05,
+        )
+
+        response_text = tokenizer.decode(outputs[0][input_ids['input_ids'].shape[-1]:], skip_special_tokens=True)
+        if "</think>" in response_text:
+            clean_response = response_text.split("</think>")[-1].strip()
+        else:
+            clean_response = response_text.strip()
+
+        print(f"Assistant: {clean_response}")
 
-outputs = model.generate(**inputs, max_new_tokens=512)
-response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        messages.append({"role": "assistant", "content": clean_response})
+    except KeyboardInterrupt:
+        print("\n再见！")
+        break
 
-print(response)
 ```
 
 ## 限制与偏见