Update app.py

app.py

@@ -5,12 +5,10 @@ from policy import load_policy_model, save_checkpoint
 from reward_fn import reward_fn
 from grpo_train import grpo_step
 
-
-# Load models
-policy_model, gen_model, ref_model, tokenizer = load_policy_model()
+policy_model, gen_model, tokenizer = load_policy_model()
 
 reward_history = []
-global_step = 0  # simple in‑memory step counter
+global_step = 0
 
 
 def plot_rewards(history):
@@ -29,7 +27,6 @@ def run_step(prompt):
     result = grpo_step(
         policy_model=policy_model,
         gen_model=gen_model,
-        ref_model=ref_model,
         tokenizer=tokenizer,
         prompt=prompt,
         reward_fn=reward_fn,
@@ -38,8 +35,8 @@ def run_step(prompt):
     reward_history.append(float(result["reward"]))
     reward_plot = plot_rewards(reward_history)
 
-    # Save checkpoint 
-    save_checkpoint(policy_model, global_step)
+    if global_step % 10 == 0:
+        save_checkpoint(policy_model, global_step)
 
     return result["text"], result["reward"], result["kl"], result["loss"], reward_plot