prm

README.md

@@ -29,19 +29,19 @@ device = 0 if torch.cuda.is_available() else "cpu"
 
 # Model initialization
 prm_model = PRMModel(
-    model_name="lvwerra/gpt2-imdb-pos-v2",
-    ref_model_name="lvwerra/gpt2-imdb",
-    reward_model_name="lvwerra/distilbert-imdb",
-    device=device,
+   model_name="lvwerra/gpt2-imdb-pos-v2",
+   ref_model_name="lvwerra/gpt2-imdb",
+   reward_model_name="lvwerra/distilbert-imdb",
+   device=device,
 )
 
 # Generation arguments
 gen_kwargs = {
-    "min_length": -1,
-    "top_k": 0.0,
-    "top_p": 1.0,
-    "do_sample": True,
-    "pad_token_id": prm_model.tokenizer.eos_token_id,
+   "min_length": -1,
+   "top_k": 0.0,
+   "top_p": 1.0,
+   "do_sample": True,
+   "pad_token_id": prm_model.tokenizer.eos_token_id,
 }
 sent_kwargs = {"top_k": None, "function_to_apply": "none", "batch_size": 16}
 

process_supervision/prm.py

@@ -4,7 +4,7 @@
 
 
 class PRMModel:
-    def __init__(self, model_name, ref_model_name, reward_model_name, device):
+    def __init__(self, model_name: str = "lvwerra/gpt2-imdb-pos-v2", ref_model_name: str = "lvwerra/gpt2-imdb", reward_model_name: str = "lvwerra/distilbert-imdb", device):
         """
         Initialize the PRM model with specified models and tokenizer.
 
@@ -17,9 +17,11 @@ def __init__(self, model_name, ref_model_name, reward_model_name, device):
         self.model = AutoModelForCausalLMWithValueHead.from_pretrained(
             model_name
         ).to(device)
+
         self.ref_model = AutoModelForCausalLMWithValueHead.from_pretrained(
             ref_model_name
         ).to(device)
+
         self.reward_pipe = pipeline(
             "sentiment-analysis", model=reward_model_name, device=device
         )