Update services/strategy.py

services/strategy.py

@@ -52,7 +52,7 @@ class BestOfN(GenerationStrategy):
                 input_ids = generator.tokenizer(prompt, return_tensors="pt").input_ids.to(self.device)
                 output = generator.generate(input_ids, **model_kwargs)
                 response =generator.tokenizer.decode(output[0], skip_special_tokens=True)
-                score = generator.prm_model(**self.llama_tokenizer(response, return_tensors="pt").to(self.device)).logits.mean().item()
+                score = generator.prm_model(**generator.tokenizer(response, return_tensors="pt").to(self.device)).logits.mean().item()
                 scored_outputs.append((response, score))
             return max(scored_outputs, key=lambda x: x[1])[0]
 
@@ -76,7 +76,7 @@ class DVT(GenerationStrategy):
                 input_ids = generator.tokenizer(prompt, return_tensors="pt").input_ids.to(self.device)
                 output = generator.generate(input_ids, **model_kwargs)
                 response = generator.tokenizer.decode(output[0], skip_special_tokens=True)
-                score = generator.prm_model(**self.llama_tokenizer(response, return_tensors="pt").to(self.device)).logits.mean().item()
+                score = generator.prm_model(**generator.tokenizer(response, return_tensors="pt").to(self.device)).logits.mean().item()
                 results.append((response, score))
             
             for _ in range(depth - 1):
@@ -85,7 +85,7 @@ class DVT(GenerationStrategy):
                     input_ids = generator.tokenizer(response, return_tensors="pt").input_ids.to(self.device)
                     output = generator.generate(input_ids, **model_kwargs)
                     extended_response = generator.tokenizer.decode(output[0], skip_special_tokens=True)
-                    score = generator.prm_model(**self.llama_tokenizer(extended_response, return_tensors="pt").to(self.device)).logits.mean().item()
+                    score = generator.prm_model(**generator.tokenizer(extended_response, return_tensors="pt").to(self.device)).logits.mean().item()
                     results.append((extended_response, score))
             return max(results, key=lambda x: x[1])[0]