Update services/strategy.py

services/strategy.py

@@ -35,21 +35,21 @@ class DefaultStrategy(GenerationStrategy):
         return generator.tokenizer.decode(output[0], skip_special_tokens=True)
     #def generate(self, generator: 'BaseGenerator', prompt: str, model_kwargs: Dict[str, Any], **kwargs) -> str:
     #    
-    #    tokenizer = generator.tokenizers["llama"]
-    #    model = generator.models["llama"].generate
+    #    tokenizer = generator.tokenizer
+    #    model = generator.model.generate
    # 
-   #     input_ids = generator.tokenizers["llama"](prompt, return_tensors="pt").input_ids.to(generator.device)
-   #     output = generator.models["llama"].generate(input_ids, **model_kwargs)
-   #     return generator.tokenizers["llama"].decode(output[0], skip_special_tokens=True)
+   #     input_ids = generator.tokenizer(prompt, return_tensors="pt").input_ids.to(generator.device)
+   #     output = generator.model.generate(input_ids, **model_kwargs)
+   #     return generator.tokenizer.decode(output[0], skip_special_tokens=True)
 
 
 class MajorityVotingStrategy(GenerationStrategy):
     def generate(self, generator: 'BaseGenerator', prompt: str, model_kwargs: Dict[str, Any], num_samples: int = 5, **kwargs) -> str:
         outputs = []
         for _ in range(num_samples):
-            input_ids = generator.tokenizers["llama"](prompt, return_tensors="pt").input_ids.to(generator.device)
-            output = generator.models["llama"].generate(input_ids, **model_kwargs)
-            outputs.append(generator.tokenizers["llama"].decode(output[0], skip_special_tokens=True))
+            input_ids = generator.tokenizer(prompt, return_tensors="pt").input_ids.to(generator.device)
+            output = generator.model.generate(input_ids, **model_kwargs)
+            outputs.append(generator.tokenizer.decode(output[0], skip_special_tokens=True))
         return max(set(outputs), key=outputs.count)
 
 
@@ -57,43 +57,43 @@ 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.tokenizers["llama"](prompt, return_tensors="pt").input_ids.to(generator.device)
-                output = generator.models["llama"].generate(input_ids, **model_kwargs)
-                response =generator.tokenizers["llama"].decode(output[0], skip_special_tokens=True)
-                score = generator.prm_model(**generator.tokenizers["llama"](response, return_tensors="pt").to(generator.device)).logits.mean().item()
+                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(**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.tokenizers["llama"](prompt, return_tensors="pt").input_ids.to(generator.device)
-            outputs = generator.models["llama"].generate(
+            input_ids = generator.tokenizer(prompt, return_tensors="pt").input_ids.to(generator.device)
+            outputs = generator.model.generate(
                 input_ids,
                 num_beams=num_samples,
                 num_return_sequences=num_samples,
                 **model_kwargs
             )
-            return [generator.tokenizers["llama"].decode(output, skip_special_tokens=True) for output in outputs]
+            return [generator.tokenizer.decode(output, skip_special_tokens=True) for output in outputs]
 
 
 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.tokenizers["llama"](prompt, return_tensors="pt").input_ids.to(generator.device)
-                output = generator.models["llama"].generate(input_ids, **model_kwargs)
-                response = generator.tokenizers["llama"].decode(output[0], skip_special_tokens=True)
-                score = generator.prm_model(**generator.tokenizers["llama"](response, return_tensors="pt").to(generator.device)).logits.mean().item()
+                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(**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.tokenizers["llama"](response, return_tensors="pt").input_ids.to(generator.device)
-                    output = generator.models["llama"].generate(input_ids, **model_kwargs)
-                    extended_response = generator.tokenizers["llama"].decode(output[0], skip_special_tokens=True)
-                    score = generator.prm_model(**generator.tokenizers["llama"](extended_response, return_tensors="pt").to(generator.device)).logits.mean().item()
+                    input_ids = generator.tokenizer(response, return_tensors="pt").input_ids.to(generator.device)
+                    output = generator.model.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(generator.device)).logits.mean().item()
                     results.append((extended_response, score))
             return max(results, key=lambda x: x[1])[0]