Update app.py

app.py

@@ -1,92 +1,84 @@
 import os
 import subprocess
-import sys
-import json
-import numpy as np
-# 设置环境
-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")
-    else:
-        repo_path = os.path.abspath("skywork-o1-prm-inference")
+import gradio as gr
+import torch
+from transformers import AutoTokenizer
 
-    if repo_path not in sys.path:
-        sys.path.append(repo_path)
-        print(f"Added {repo_path} to Python path")
+# 克隆项目仓库
+REPO_URL = "https://github.com/SkyworkAI/skywork-o1-prm-inference.git"
+REPO_DIR = "skywork-o1-prm-inference"
 
-setup_environment()
+# 在 Hugging Face Spaces 上首次启动时克隆
+if not os.path.exists(REPO_DIR):
+    subprocess.run(["git", "clone", REPO_URL], check=True)
 
-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
+# 导入必要的模块
+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
+)
 
-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()
+# 模型配置
+MODEL_ID = "Skywork/Skywork-o1-Open-PRM-Qwen-2.5-1.5B"
 
-def evaluate(problem, response):
-    try:
-        # 处理输入数据
-        processed_data = prepare_input(problem, response, tokenizer=tokenizer, step_token="\n")
-        input_ids = [processed_data[0]]  # 第一个元素
-        steps = [processed_data[1]]      # 第二个元素
-        reward_flags = [processed_data[2]]  # 第三个元素
+# 初始化模型和tokenizer
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+model = PRM_MODEL.from_pretrained(MODEL_ID).to("cpu").eval()
 
-        # 准备批处理输入
-        input_ids, attention_mask, reward_flags = prepare_batch_input_for_model(
-            input_ids,
-            reward_flags, 
-            tokenizer.pad_token_id
-        )
+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
+
+    # 转换为批处理形式
+    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")
 
-        # 确保在CPU上
-        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
+        )
 
-        # 模型推理
-        with torch.no_grad():
-            _, _, rewards = model(
-                input_ids=input_ids, 
-                attention_mask=attention_mask, 
-                return_probs=True
-            )
+    # 获取step rewards
+    step_rewards = derive_step_rewards(rewards, reward_flags)
+    return step_rewards[0]
 
-        # 计算步骤奖励
-        step_rewards = derive_step_rewards(rewards, reward_flags)
+def inference_interface(problem, response):
+    rewards = compute_step_rewards(problem, response)
+    return rewards
 
-        # 确保返回的是有效的JSON字符串
-        if isinstance(step_rewards[0], torch.Tensor):
-            return json.dumps(step_rewards[0].cpu().numpy().tolist())
-        elif isinstance(step_rewards[0], np.ndarray):
-            return json.dumps(step_rewards[0].tolist())
-        else:
-            return json.dumps(list(step_rewards[0]))  # 转换为列表
-    except Exception as e:
-        return json.dumps({"error": str(e)})
-# 创建Gradio界面
-iface = gr.Interface(
-    fn=evaluate,
+# Gradio界面配置
+interface = gr.Interface(
+    fn=inference_interface,
     inputs=[
-        gr.Textbox(label="Problem", lines=4),
-        gr.Textbox(label="Response", lines=8)
-    ],
-    outputs=gr.JSON(),
-    title="Problem Response Evaluation",
-    description="Enter a problem and its response to get step-wise rewards",
-    examples=[
-        [
-            "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?",
-            "To determine how much money Janet makes every day at the farmers' market, we need to follow these steps:\n1. Calculate the total number of eggs laid by the ducks per day.\n   Janet's ducks lay 16 eggs per day.\n2. Determine the number of eggs Janet uses each day.\n   - She eats 3 eggs for breakfast every morning.\n   - She bakes muffins for her friends every day with 4 eggs.\n   So, the total number of eggs used per day is:\n   3 + 4 = 7 eggs\n3. Calculate the number of eggs Janet sells at the farmers' market each day.\n   Subtract the number of eggs used from the total number of eggs laid:\n   16 - 7 = 9 eggs\n4. Determine how much money Janet makes from selling the eggs.\n   She sells each egg for $2, so the total amount of money she makes is:\n   9 ×2 = 18 dollars\nTherefore, the amount of money Janet makes every day at the farmers' market is $18."
-        ]
+        gr.Textbox(lines=4, label="Problem (题目)"),
+        gr.Textbox(lines=6, label="Response (回答)"),
     ],
-    cache_examples=False  # 禁用示例缓存
+    outputs="json",
+    title="Skywork-o1-prm-inference Demo",
+    description=(
+        "输入题目和回答，点击提交查看其每个 step 对应的 reward，"
+        "这些 reward 值用于度量回答每一步的质量。"
+    ),
 )
 
-# 启动接口
-iface.launch(server_name="0.0.0.0")
+if __name__ == "__main__":
+    interface.launch(server_name="0.0.0.0", server_port=7860)
+