Create app.py

app.py

@@ -0,0 +1,39 @@
+import gradio as gr
+from transformers import AutoTokenizer
+from model_utils.prm_model import PRM_MODEL
+from model_utils.io_utils import prepare_input, prepare_batch_input_for_model, derive_step_rewards
+import torch
+
+# 初始化模型和tokenizer (和你现有代码一样)
+model_id = "Skywork/Skywork-o1-Open-PRM-Qwen-2.5-1.5B"
+tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+model = PRM_MODEL.from_pretrained(model_id).to("cpu").eval()
+
+def evaluate(problem, response):
+    processed_data = prepare_input(problem, response, tokenizer=tokenizer, step_token="\n")
+    input_ids, steps, reward_flags = [processed_data]
+    
+    input_ids, attention_mask, reward_flags = prepare_batch_input_for_model(input_ids, reward_flags, tokenizer.pad_token_id)
+    input_ids = input_ids.to("cpu")
+    attention_mask = attention_mask.to("cpu")
+    
+    with torch.no_grad():
+        _, _, rewards = model(input_ids=input_ids, attention_mask=attention_mask, return_probs=True)
+    
+    step_rewards = derive_step_rewards(rewards, reward_flags)
+    return step_rewards[0].tolist()
+
+# 创建Gradio界面
+iface = gr.Interface(
+    fn=evaluate,
+    inputs=[
+        gr.Textbox(label="Problem"),
+        gr.Textbox(label="Response")
+    ],
+    outputs=gr.JSON(label="Step Rewards"),
+    title="Problem Response Evaluation",
+    description="Enter a problem and its response to get step-wise rewards"
+)
+
+# 启动接口
+iface.launch()