Update inference.py

inference.py

@@ -8,6 +8,8 @@ peft_model_id = "philschmid/gemma-7b-dolly-chatml"
 tokenizer = AutoTokenizer.from_pretrained(peft_model_id)
 model = AutoPeftModelForCausalLM.from_pretrained(peft_model_id, device_map="auto", torch_dtype=torch.float16)
 pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
+eos_token = tokenizer("<|im_end|>",add_special_tokens=False)["input_ids"][0]
+print(f"eos_token: {eos_token}")
 
 # run inference
 messages = [
@@ -18,8 +20,10 @@ messages = [
 ]
 
 prompt = pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-outputs = pipe(prompt, max_new_tokens=1024, do_sample=True, temperature=0.7, top_k=50, top_p=0.95, pad_token_id=pipe.tokenizer.pad_token_id, eos_token_id=pipe.tokenizer.eos_token_id)
-print(outputs[0]["generated_text"])
+outputs = pipe(prompt, max_new_tokens=1024, do_sample=True, temperature=0.7, top_k=50, top_p=0.95, eos_token_id=eos_token)
+
+print(f"prompt:\n {messages[0]['content']}")
+print(f"response:\n {outputs[0]['generated_text'][len(prompt):]}")
 
 # run inference
 messages = [
@@ -30,8 +34,7 @@ messages = [
 ]
 
 prompt = pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-outputs = pipe(prompt, max_new_tokens=1024, do_sample=True, temperature=0.7, top_k=50, top_p=0.95, pad_token_id=pipe.tokenizer.pad_token_id, eos_token_id=pipe.tokenizer.eos_token_id)
-print(outputs[0]["generated_text"])
-
+outputs = pipe(prompt, max_new_tokens=1024, do_sample=True, temperature=0.7, top_k=50, top_p=0.95, eos_token_id=eos_token)
 
-# pip3 list | grep -e transformers -e peft -e torch -e trl -e accelerate
+print(f"prompt:\n {messages[0]['content']}")
+print(f"response:\n {outputs[0]['generated_text'][len(prompt):]}")