update

examples/tutorials/rlhf/gpt2_sst2_ppo/step_3_generation.py

@@ -1,7 +1,23 @@
 #!/usr/bin/python3
 # -*- coding: utf-8 -*-
+"""
+python3 step_3_generation.py
+
+"""
 import argparse
-import time
+import os
+from pathlib import Path
+import platform
+import sys
+
+# 路径配置
+if platform.system() in ("Windows", "Darwin"):
+    from project_settings import project_path, temp_directory
+else:
+    project_path = Path(os.path.abspath("../../../../"))
+    temp_directory = Path("/root/autodl-tmp/OpenMiniMind/temp")
+
+sys.path.append(project_path.as_posix())
 
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
@@ -13,7 +29,7 @@ def get_args():
     parser = argparse.ArgumentParser()
     parser.add_argument(
         "--pretrained_model_name_or_path",
-        default=(project_path / "trained_models/gpt2-sst2-ppo/checkpoint-1250"),
+        default=(project_path / "trained_models/gpt2-sst2-ppo/checkpoint-50"),
         # default=(project_path / "trained_models/gpt2-sst2-ppo/checkpoint-150"),
         # default=(project_path / "trained_models/gpt2-sst2-generation-epoch-3/checkpoint-5400"),
         type=str