Update app.py

app.py

@@ -7,11 +7,9 @@ import json
 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")
+        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")
+        sys.path.append(repo_path)print(f"Added {repo_path} to Python path")
 
 setup_environment()
 
@@ -27,26 +25,28 @@ model = PRM_MODEL.from_pretrained(model_id).to("cpu").eval()
 
 def evaluate(problem, response):
     try:
+        # 处理输入数据
         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, 
+        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
-        )
-        
+        )# 确保在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
             )
-        step_rewards = derive_step_rewards(rewards, reward_flags)
-        #确保返回的是有效的JSON字符串
+        # 计算步骤奖励
+        step_rewards = derive_step_rewards(rewards, reward_flags)# 确保返回的是有效的JSON字符串
         return json.dumps(step_rewards[0].tolist())
     except Exception as e:
         return json.dumps({"error": str(e)})
@@ -63,11 +63,11 @@ iface = gr.Interface(
     description="Enter a problem and its response to get step-wise rewards",
     examples=[
         [
-            "Janet'sducks lay 16 eggs per day...",
+            "Janet's ducks lay 16 eggs per day...",
             "To determine how much money Janet makes..."
         ]
     ],
-    cache_examples=False# 禁用示例缓存
+    cache_examples=False  # 禁用示例缓存
 )
 
 # 启动接口