Update services/strategy.py

services/strategy.py

@@ -66,12 +66,15 @@ class BestOfN(GenerationStrategy):
             response = generator.tokenizer.decode(output[0], skip_special_tokens=True)
 
             # Tokenize the response for scoring with the PRM model
-            #TODO use the real tokenizer from generator
             response_inputs = generator.tokenizer(response, return_tensors="pt").to(generator.device)
 
-            # Pass the response inputs correctly to the PRM model
-            prm_output = generator.prm_model(response_inputs)  # Pass the inputs correctly to the model
+            # Extract the necessary inputs for prm_model
+            prm_input_ids = response_inputs["input_ids"]  # Always present
+            attention_mask = response_inputs["attention_mask"]  # Optional, depending on your model
             
+            # Pass only the required tensors to prm_model
+            prm_output = generator.prm_model(input_ids=prm_input_ids, attention_mask=attention_mask)
+ 
             # Check the expected output structure for prm_model and use it accordingly
             score = prm_output.logits.mean().item() if hasattr(prm_output, 'logits') else 0.0