Update app.py

app.py

@@ -1,84 +1,65 @@
 import os
 import subprocess
-import gradio as gr
+import sys
 import torch
+import gradio as gr
 from transformers import AutoTokenizer
 
-# 克隆项目仓库
-REPO_URL = "https://github.com/SkyworkAI/skywork-o1-prm-inference.git"
-REPO_DIR = "skywork-o1-prm-inference"
+# 设置环境：若本地无skywork-o1-prm-inference目录，则git clone下来，并将其加入sys.path
+def setup_environment():
+    if not os.path.exists("skywork-o1-prm-inference"):
+        print("Cloning repository...")
+        subprocess.run(["git", "clone", "https://github.com/SkyworkAI/skywork-o1-prm-inference.git"], check=True)
+    repo_path = os.path.abspath("skywork-o1-prm-inference")
+    if repo_path not in sys.path:
+        sys.path.append(repo_path)
+        print(f"Added {repo_path} to Python path")
 
-# 在 Hugging Face Spaces 上首次启动时克隆
-if not os.path.exists(REPO_DIR):
-    subprocess.run(["git", "clone", REPO_URL], check=True)
+setup_environment()
 
-# 导入必要的模块
+# 在环境设置后再导入项目内的模块
 from skywork-o1-prm-inference.model_utils.prm_model import PRM_MODEL
-from skywork-o1-prm-inference.model_utils.io_utils import (
-    prepare_input,
-    prepare_batch_input_for_model,
-    derive_step_rewards
-)
+from skywork-o1-prm-inference.model_utils.io_utils import prepare_input, prepare_batch_input_for_model, derive_step_rewards
 
-# 模型配置
-MODEL_ID = "Skywork/Skywork-o1-Open-PRM-Qwen-2.5-1.5B"
+# 模型ID
+model_id = "Skywork/Skywork-o1-Open-PRM-Qwen-2.5-1.5B"
 
 # 初始化模型和tokenizer
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
-model = PRM_MODEL.from_pretrained(MODEL_ID).to("cpu").eval()
+tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+model = PRM_MODEL.from_pretrained(model_id).to("cpu").eval()
 
-def compute_step_rewards(problem_text, response_text):
-    # 准备输入数据
-    processed_input = prepare_input(
-        problem_text,
-        response_text,
-        tokenizer=tokenizer,
-        step_token="\n"
-    )
-    input_ids, steps, reward_flags = processed_input
+def compute_rewards(problem, response):
+    # 准备数据
+    data = {
+        "problem": problem,
+        "response": response
+    }
+    processed_data = [prepare_input(data["problem"], data["response"], tokenizer=tokenizer, step_token="\n")]
+    input_ids, steps, reward_flags = zip(*processed_data)
 
-    # 转换为批处理形式
-    input_ids, attention_mask, reward_flags = prepare_batch_input_for_model(
-        [input_ids],
-        [reward_flags],
-        tokenizer.pad_token_id
-    )
+    # 准备 batch 输入
+    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")
     if isinstance(reward_flags, torch.Tensor):
         reward_flags = reward_flags.to("cpu")
 
-    # 模型推理
     with torch.no_grad():
-        _, _, rewards = model(
-            input_ids=input_ids, 
-            attention_mask=attention_mask, 
-            return_probs=True
-        )
+        _, _, rewards = model(input_ids=input_ids, attention_mask=attention_mask, return_probs=True)
 
-    # 获取step rewards
+    # 计算 step_rewards
     step_rewards = derive_step_rewards(rewards, reward_flags)
     return step_rewards[0]
 
-def inference_interface(problem, response):
-    rewards = compute_step_rewards(problem, response)
-    return rewards
+# 搭建 Gradio 界面
+with gr.Blocks() as demo:
+    gr.Markdown("# PRM Reward Calculation")
+    problem_input = gr.Textbox(label="Problem", lines=5, value="Janet's ducks lay 16 eggs per day. She eats three for breakfast every morning and bakes muffins for her friends every day with four. She sells the remainder at the farmers' market daily for $2 per fresh duck egg. How much in dollars does she make every day at the farmers' market?")
+    response_input = gr.Textbox(label="Response", lines=10, value="To determine how much money Janet makes every day at the farmers' market, we need to follow these steps:\n1. ...")
+    output = gr.JSON(label="Step Rewards")
 
-# Gradio界面配置
-interface = gr.Interface(
-    fn=inference_interface,
-    inputs=[
-        gr.Textbox(lines=4, label="Problem (题目)"),
-        gr.Textbox(lines=6, label="Response (回答)"),
-    ],
-    outputs="json",
-    title="Skywork-o1-prm-inference Demo",
-    description=(
-        "输入题目和回答，点击提交查看其每个 step 对应的 reward，"
-        "这些 reward 值用于度量回答每一步的质量。"
-    ),
-)
+    submit_btn = gr.Button("Compute Rewards")
+    submit_btn.click(fn=compute_rewards, inputs=[problem_input, response_input], outputs=output)
 
 if __name__ == "__main__":
-    interface.launch(server_name="0.0.0.0", server_port=7860)
-
+    demo.launch(server_name="0.0.0.0", server_port=7860)