Update services/strategy.py

services/strategy.py

@@ -49,17 +49,17 @@ 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(self.device)
+                input_ids = generator.tokenizer(prompt, return_tensors="pt").input_ids.to(generator.device)
                 output = generator.generate(input_ids, **model_kwargs)
                 response =generator.tokenizer.decode(output[0], skip_special_tokens=True)
-                score = generator.prm_model(**generator.tokenizer(response, return_tensors="pt").to(self.device)).logits.mean().item()
+                score = generator.prm_model(**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]
 
 
 class BeamSearch(GenerationStrategy):
     def generate(self, generator: 'BaseGenerator', prompt: str, model_kwargs: Dict[str, Any], num_samples: int = 5, **kwargs) -> str:
-            input_ids = generator.tokenizer(prompt, return_tensors="pt").input_ids.to(self.device)
+            input_ids = generator.tokenizer(prompt, return_tensors="pt").input_ids.to(generator.device)
             outputs = generator.generate(
                 input_ids,
                 num_beams=num_samples,
@@ -73,19 +73,19 @@ class DVT(GenerationStrategy):
     def generate(self, generator: 'BaseGenerator', prompt: str, model_kwargs: Dict[str, Any], num_samples: int = 5, **kwargs) -> str:
             results = []
             for _ in range(breadth):
-                input_ids = generator.tokenizer(prompt, return_tensors="pt").input_ids.to(self.device)
+                input_ids = generator.tokenizer(prompt, return_tensors="pt").input_ids.to(generator.device)
                 output = generator.generate(input_ids, **model_kwargs)
                 response = generator.tokenizer.decode(output[0], skip_special_tokens=True)
-                score = generator.prm_model(**generator.tokenizer(response, return_tensors="pt").to(self.device)).logits.mean().item()
+                score = generator.prm_model(**generator.tokenizer(response, return_tensors="pt").to(generator.device)).logits.mean().item()
                 results.append((response, score))
             
             for _ in range(depth - 1):
                 best_responses = sorted(results, key=lambda x: x[1], reverse=True)[:breadth]
                 for response, _ in best_responses:
-                    input_ids = generator.tokenizer(response, return_tensors="pt").input_ids.to(self.device)
+                    input_ids = generator.tokenizer(response, return_tensors="pt").input_ids.to(generator.device)
                     output = generator.generate(input_ids, **model_kwargs)
                     extended_response = generator.tokenizer.decode(output[0], skip_special_tokens=True)
-                    score = generator.prm_model(**generator.tokenizer(extended_response, return_tensors="pt").to(self.device)).logits.mean().item()
+                    score = generator.prm_model(**generator.tokenizer(extended_response, return_tensors="pt").to(generator.device)).logits.mean().item()
                     results.append((extended_response, score))
             return max(results, key=lambda x: x[1])[0]