Update services/strategy.py

services/strategy.py

@@ -55,14 +55,25 @@ class MajorityVotingStrategy(GenerationStrategy):
 
 class BestOfN(GenerationStrategy):
     def generate(self, generator: 'BaseGenerator', prompt: str, model_kwargs: Dict[str, Any], num_samples: int = 5, **kwargs) -> str:
-            scored_outputs = []
-            for _ in range(num_samples):
-                input_ids = generator.tokenizer(prompt, return_tensors="pt").input_ids.to(generator.device)
-                output = generator.model.generate(input_ids, **model_kwargs)
-                response =generator.tokenizer.decode(output[0], skip_special_tokens=True)
-                score = generator.prm_model.generate(**generator.tokenizer(response, return_tensors="pt").to(generator.device)).logits.mean().item()
-                scored_outputs.append((response, score))
-            return max(scored_outputs, key=lambda x: x[1])[0]
+        scored_outputs = []
+        for _ in range(num_samples):
+            # Tokenize the prompt and move tensors to the appropriate device
+            input_ids = generator.tokenizer(prompt, return_tensors="pt").input_ids.to(generator.device)
+
+            # Generate output from the main model
+            output = generator.model.generate(input_ids, **model_kwargs)
+            response = generator.tokenizer.decode(output[0], skip_special_tokens=True)
+
+            # Tokenize the response for scoring with the PRM model
+            response_inputs = generator.tokenizer(response, return_tensors="pt").to(generator.device)
+            prm_output = generator.prm_model(**response_inputs)  # Pass the inputs correctly to the model
+            score = prm_output.logits.mean().item()
+
+            # Append the response and its score
+            scored_outputs.append((response, score))
+
+        # Return the response with the highest score
+        return max(scored_outputs, key=lambda x: x[1])[0]
 
 
 class BeamSearch(GenerationStrategy):