Update README.md

README.md

@@ -37,6 +37,94 @@ Perform rejection sampling for each steps until the end.
 |--|--|
 |![sequential_rejection_sampling_zh_prompt.png](image%2Fsequential_rejection_sampling_zh_prompt.png)|![sequential_rejection_sampling_zh.png](image%2Fsequential_rejection_sampling_zh.png)|
 
+### Basic Usage
+
+```
+from transformers import pipeline
+
+# load reward model
+pipe = pipeline("token-classification", model="mrzjy/NovelWriting-Outline-PRM-Qwen2.5-0.5B-Reward", device="cuda")
+separator = "\n"  # It's important to use the same separator as the one used during training
+
+def evaluate_reward(batch_prompt, batch_steps) -> dict:
+    # Add a separator between the prompt and each steps
+    rewards = []
+    assert len(batch_prompt) == len(batch_steps)
+    for prompt, steps in zip(batch_prompt, batch_steps):
+        reward_dict = []
+        for idx in range(1, len(steps) + 1):
+            current_steps = steps[0:idx]
+            text = separator.join((prompt, *current_steps)) + separator  # Add a separator between the prompt and each steps
+            pred = pipe(text)[-1]
+            score, pred_entity = pred["score"], pred["entity"]
+            # always use the score for positive label
+            if pred_entity == "LABEL_0":
+                score = 1 - score
+            reward_dict.append({"step": steps[idx-1], "score": score})
+        rewards.append(reward_dict)
+    return rewards
+
+
+batch_prompt = [
+    """你是一位专业小说写手。你的任务是基于以下故事灵感进行小说大纲创作。\n一个关于李向东穿越回1979年，面对家庭压力和挑战，最终与家人共度温馨时光的故事。\n\n以下是你设计的角色：\n：\n\n角色1：李向东，主角，年轻时下乡插队，现在回城后无所事事，渴望改变现状。\n角色2：李母，配角，李向东的母亲，唠叨且疼爱儿子，希望儿子能有份好工作。\n角色3：李父，配角，李向东的父亲，为人温和，对儿子寄予厚望。\n角色4：周玉琴，配角，李向东的妻子，性格温和，对丈夫和孩子充满关爱。\n角色5：李晓江，配角，李向东的大侄子，性格顽皮，喜欢和三叔玩闹。\n角色6：李晓涛，配角，李向东的二侄子，性格憨厚，喜欢和三叔一起干活。\n角色7：李晓梅，配角，李向东的侄女，活泼好动，喜欢和哥哥姐姐一起玩耍。\n角色8：李晓兰，配角，李向东的侄女，与妹妹一起，喜欢和哥哥姐姐一起玩耍。\n角色9：李晓海，配角，李向东的儿子，年仅三岁，性格活泼，喜欢和哥哥姐姐一起玩耍。\n角色10：张姨，配角，李母的老姐妹，经常在供销社上班，对李向东关心有加。\n角色11：李卫国，配角，李向东的大哥，为人老实，喜欢和三叔一起干活。\n角色12：李卫民，配角，李向东的二哥，为人憨厚，喜欢和三叔一起干活。\n角色13：李小竹，配角，李向东的女儿，活泼好动，喜欢和哥哥姐姐一起玩耍。\n角色14：刘大娘，配角，住在隔壁的邻居，经常和李母一起打叶子牌。\n\n\n请基于以上信息为小说前10章设计故事大纲，每个大纲大概在65字左右。""",
+    """Act as a novel writer. Your task is to craft novel outlines based on the following story idea:\nA story about Apari, who awakens in a serene room after being trapped in an underground shelter during a violent storm. The stark contrast between the harsh outside and the cozy inside heightens his confusion and unease, as he grapples with recent vows and the nature of his current situation.\n\nHere's your character design:\nApari, supporting, A young man hiding in an underground shelter, screaming in fear, his presence unknown to the reader.\nNowra, supporting, A companion silent in the narrative, her absence a mystery to Apari.\nTheo, protagonist, A man sharing a room with the narrator, his face serious and changed, unsure if their vow to the moon goddess will be accepted.\nIsland Resident, supporting, An anonymous figure, possibly a guard, whose disappearance adds to the tension, unseen but felt.\n\nPlease create the story outlines for the first 5 chapters, with each chapter outline in 65 words."""
+]
+
+batch_steps = [
+    """第1章：李向东重返1979年，面对母亲的责备与邻居的流言蜚语，他内心回想家族往事，努力调整心态，试图摆脱母亲喋喋不休的困扰，最后以去厕所为由脱身。
+第2章：李向东在胡同里闲逛，偶遇妻子周玉琴带着女儿玩耍，周玉琴提醒他回家吃饭。李向东将女儿带回家，通过与女儿的互动，表达对家庭的责任感。
+第3章：在共进晚餐时，家庭氛围略显紧张，父亲批评他无所作为，母亲则坚定地支持他。李向东感受着来自家庭的压力，决心改变现状。父亲则要求他去街道办事处报到。
+第4章：在一次家庭会议上，李向东与家人讨论工作机会，得知火车站正在招聘。母亲担心老人健康，早早起来买肉。李向东关心奶奶，��过装满肉的布袋，奶奶感激地留下煮好的鸡蛋给孙子。
+第5章：李向东回家探望家人，被母亲发现偷吃鸡蛋，引发了一次充满亲情的对话。李老太讲述偏爱李向东的缘由，并以过去煤票的故事教育孙子，展现了家庭的温情和长辈的智慧。
+第6章：李向东在家中发现了祖传的明朝红酸枝八仙桌，得知其来历后，他决定将桌子搬回西厢房，却被爷爷质疑其价值。
+第7章：李向东发现桌子背面沾有鼻涕，责问孩子们后，许诺给他们奶油雪糕，促使孩子们帮忙擦拭干净，整个过程中，李向东与孩子们互动，展现了家庭的温暖与幽默。
+第8章：在一个夏日的午后，李向东带着孩子们去买冰棍。在买冰棍的过程中，他与大侄子李晓江交谈，了解他成长过程中的情感世界。
+第9章：李向东回城后，首次前往供销社买冰棍，受到了张姨的关心和唠叨。回到家后，他将冰棍分给孩子们，却引发了母亲的责备和担忧。
+第10章：李向东为爷爷奶奶买了冰棍，却发现儿子因为惹祸被罚站。他与妻子分享冰棍，孩子们争抢剩下的小木棒，为的是换取冰棍。""".splitlines(),
+    """Chapter 1: Strong winds, heavy rain, and crashing waves set a gloomy and foreboding atmosphere. Overcast skies and biting cold in February create an eerie mood, with clouds blanketing the island. The natural elements hint at an impending storm and an ominous feeling that hangs over the scene.
+Chapter 2: Trapped in an underground shelter, Apari's screams echo in the damp, echoing space. A violent storm rages outside, with howling winds and torrential rain battering the island. Thoughts of the guards and the absence of Nowra's comforting voice heighten his sense of dread, leaving him isolated and vulnerable in the midst of chaos.
+Chapter 3: Apari awakens in a starkly contrasting environment, surrounded by the comforting scent of clean sheets and lavender. The soft, cozy bed contrasts with the harsh reality of his previous isolation. He sits up, the blankets pooling around him, to find a sleeping man beside him, half-naked and peaceful. This serene and unexpected intimacy heightens Apari’s confusion and unease, leaving him to wonder about the circumstances that brought him here and the nature of the person sharing the bed.
+Chapter 4: Staring at the face of the man beside him, Apari notes the seriousness in his expression. Doubts arise about the recent vow made, wondering if the moon goddess will accept it. The stark contrast between the harsh, stormy environment outside and the serene, intimate setting inside adds to Apari’s confusion and unease, leaving him questioning the nature of their current situation and the reliability of the vow.
+Chapter 5: Theo's gentle, chaste kiss introduces a moment of tender intimacy amidst the stark contrast between the ominous storm outside and the serene, intimate setting of their room. Apari's confusion and unease deepen as he grapples with the recent vow and the moon goddess's acceptance. The sudden shift in atmosphere heightens the sense of mystery and tension, leaving Apari questioning the nature of their current situation and the reliability of their surroundings.""".splitlines()
+]
+
+rewards = evaluate_reward(batch_prompt, batch_steps)
+```
+
+- Output should be as follows:
+
+```
+[[{'step': '第1章：李向东重返1979年，面对母亲的责备与邻居的流言蜚语，他内心回想家族往事，努力调整心态，试图摆脱母亲喋喋不休的困扰，最后以去厕所为由脱身。',
+   'score': 1.0},
+  {'step': '第2章：李向东在胡同里闲逛，偶遇妻子周玉琴带着女儿玩耍，周玉琴提醒他回家吃饭。李向东将女儿带回家，通过与女儿的互动，表达对家庭的责任感。',
+   'score': 0.99999845},
+  {'step': '第3章：在共进晚餐时，家庭氛围略显紧张，父亲批评他无所作为，母亲则坚定地支持他。李向东感受着来自家庭的压力，决心改变现状。父亲则要求他去街道办事处报到。',
+   'score': 0.99999905},
+  {'step': '第4章：在一次家庭会议上，李向东与家人讨论工作机会，得知火车站正在招聘。母亲担心老人健康，早早起来买肉。李向东关心奶奶，接过装满肉的布袋，奶奶感激地留下煮好的鸡蛋给孙子。',
+   'score': 1.0},
+  {'step': '第5章：李向东回家探望家人，被母亲发现偷吃鸡蛋，引发了一次充满亲情的对话。李老太讲述偏爱李向东的缘由，并以过去煤票的故事教育孙子，展现了家庭的温情和长辈的智慧。',
+   'score': 1.0},
+  {'step': '第6章：李向东在家中发现了祖传的明朝红酸枝八仙桌，得知其来历后，他决定将桌子搬回西厢房，却被爷爷质疑其价值。',
+   'score': 1.0},
+  {'step': '第7章：李向东发现桌子背面沾有鼻涕，责问孩子们后，许诺给他们奶油雪糕，促使孩子们帮忙擦拭干净，整个过程中，李���东与孩子们互动，展现了家庭的温暖与幽默。',
+   'score': 0.9999999},
+  {'step': '第8章：在一个夏日的午后，李向东带着孩子们去买冰棍。在买冰棍的过程中，他与大侄子李晓江交谈，了解他成长过程中的情感世界。',
+   'score': 0.03910404443740845},
+  {'step': '第9章：李向东回城后，首次前往供销社买冰棍，受到了张姨的关心和唠叨。回到家后，他将冰棍分给孩子们，却引发了母亲的责备和担忧。',
+   'score': 0.99996746},
+  {'step': '第10章：李向东为爷爷奶奶买了冰棍，却发现儿子因为惹祸被罚站。他与妻子分享冰棍，孩子们争抢剩下的小木棒，为的是换取冰棍。',
+   'score': 1.0}],
+ [{'step': 'Chapter 1: Strong winds, heavy rain, and crashing waves set a gloomy and foreboding atmosphere. Overcast skies and biting cold in February create an eerie mood, with clouds blanketing the island. The natural elements hint at an impending storm and an ominous feeling that hangs over the scene.',
+   'score': 1.0},
+  {'step': "Chapter 2: Trapped in an underground shelter, Apari's screams echo in the damp, echoing space. A violent storm rages outside, with howling winds and torrential rain battering the island. Thoughts of the guards and the absence of Nowra's comforting voice heighten his sense of dread, leaving him isolated and vulnerable in the midst of chaos.",
+   'score': 0.99999964},
+  {'step': 'Chapter 3: Apari awakens in a starkly contrasting environment, surrounded by the comforting scent of clean sheets and lavender. The soft, cozy bed contrasts with the harsh reality of his previous isolation. He sits up, the blankets pooling around him, to find a sleeping man beside him, half-naked and peaceful. This serene and unexpected intimacy heightens Apari’s confusion and unease, leaving him to wonder about the circumstances that brought him here and the nature of the person sharing the bed.',
+   'score': 1.0},
+  {'step': 'Chapter 4: Staring at the face of the man beside him, Apari notes the seriousness in his expression. Doubts arise about the recent vow made, wondering if the moon goddess will accept it. The stark contrast between the harsh, stormy environment outside and the serene, intimate setting inside adds to Apari’s confusion and unease, leaving him questioning the nature of their current situation and the reliability of the vow.',
+   'score': 0.99999964},
+  {'step': "Chapter 5: Theo's gentle, chaste kiss introduces a moment of tender intimacy amidst the stark contrast between the ominous storm outside and the serene, intimate setting of their room. Apari's confusion and unease deepen as he grapples with the recent vow and the moon goddess's acceptance. The sudden shift in atmosphere heightens the sense of mystery and tension, leaving Apari questioning the nature of their current situation and the reliability of their surroundings.",
+   'score': 1.0}]]
+```
 
 ## 1. Task Definition