from transformers import AutoTokenizer, AutoModelForCausalLM
model_id = "wangrongsheng/DPDG-Llama-8B-qlora"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
model_id, torch_dtype="auto", device_map="auto"
)
prompt = "\nPlease generate two different quality responses based on the given prompt. The first response should be a high-quality, satisfactory answer, representing the \"chosen\" option. The second response should be a low-quality, less ideal answer, representing the \"rejected\" option.\n\nWhen generating these two responses, please note the following:\n\n1. The \"chosen\" response should have substantive content, fluent expression, and be able to fully answer the question or requirement posed in the prompt.\n\n2. The \"rejected\" response can have some issues, such as illogical flow, incomplete information, or unclear expression. But ensure that it is still a response that can be loosely understood, not completely irrelevant or meaningless content.\n\n3. The lengths of the two responses should be roughly comparable, not vastly different.\n\n4. Try to reflect a clear difference in quality between the \"chosen\" response and the \"rejected\" response, so that the distinction is evident.\n\nPlease generate a \"chosen\" response and a \"rejected\" response for the given prompt according to these guidelines. This will help train the reward model to distinguish high-quality and low-quality responses.\n\nThe prompt is:\nIf an electric train is traveling south, which way is the smoke going?"
messages = [
{"role": "user", "content": prompt},
]
input_ids = tokenizer.apply_chat_template(
messages, add_generation_prompt=True, return_tensors="pt"
).to(model.device)
outputs = model.generate(
input_ids,
max_new_tokens=8192,
do_sample=True,
temperature=0.6,
top_p=0.9,
)
response = outputs[0][input_ids.shape[-1]:]
print(tokenizer.decode(response, skip_special_tokens=True))
