modify: modify mistral inference code

README.md

@@ -42,12 +42,8 @@ Here provides a code snippet with `apply_chat_template` to show you how to load
 from transformers import AutoModelForCausalLM, AutoTokenizer
 device = "cuda" # the device to load the model onto
 
-model = AutoModelForCausalLM.from_pretrained(
-    "yuyouyu/Mistral-Nemo-BD-RP",
-    torch_dtype="auto",
-    device_map="auto"
-)
-tokenizer = AutoTokenizer.from_pretrained("yuyouyu/Mistral-Nemo-BD-RP")
+
+chatbot = pipeline("text-generation", model="yuyouyu/Mistral-Nemo-BD-RP", device_map="auto")
 
 system_prompt_temp = """I want you to answer questions as if you are {role_name}, assuming you live in the world of {world} and mimicking {role_name}'s personality and speaking style. Use the tone, manner, and vocabulary that {role_name} would use. Please do not reveal that you are an AI or language model; you must always remember you are {role_name}.
 {role_name}'s character traits are {character}.
@@ -87,26 +83,8 @@ messages = [
     {"role": "system", "content": system_prompt},
     {"role": "user", "content": prompt}
 ]
-text = tokenizer.apply_chat_template(
-    messages,
-    tokenize=False,
-    add_generation_prompt=True
-)
-model_inputs = tokenizer([text], return_tensors="pt").to(device)
-
-generated_ids = model.generate(
-            model_inputs.input_ids,
-            max_new_tokens=256,
-            do_sample=True,
-            temperature=0.7,
-            repetition_penalty=1.2,
-)
-
-generated_ids = [
-    output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
-]
 
-response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+response = chatbot(messages, max_new_tokens=256, pad_token_id=chatbot.tokenizer.eos_token_id, do_sample=True, temperature=0.7)[0]['generated_text'][-1]['content']
 ```
 
 > [!IMPORTANT]