Update services/strategy.py

services/strategy.py

@@ -30,75 +30,79 @@ class GenerationStrategy(ABC):
 class DefaultStrategy(GenerationStrategy):
     @observe()
     def generate(self, generator: 'BaseGenerator', prompt: str, model_kwargs: Dict[str, Any], **kwargs) -> str:
-        input_ids = generator.tokenizer(prompt, return_tensors="pt").input_ids.to(generator.device)
-        output = generator.generate(input_ids, **model_kwargs)
-        return generator.tokenizer.decode(output[0], skip_special_tokens=True)
+        
+        tokenizer = generator.tokenizers["llama"]
+        model = generator.models["llama"].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)
 
 
 class MajorityVotingStrategy(GenerationStrategy):
-    def generate(self, generator: 'BaseGenerator', prompt: str, model_kwargs: Dict[str, Any], num_samples: int = 5, **kwargs) -> str:
+    def generate(self, generator: 'BaseGenerator', prompt: str, model_kwargs: Dict[str, Any], **kwargs, num_samples: int = 5) -> str:
         outputs = []
         for _ in range(num_samples):
-            input_ids = generator.tokenizer(prompt, return_tensors="pt").input_ids.to(generator.device)
-            output = generator.generate(input_ids, **model_kwargs)
-            outputs.append(generator.tokenizer.decode(output[0], skip_special_tokens=True))
+            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))
         return max(set(outputs), key=outputs.count)
 
 
 class BestOfN(GenerationStrategy):
-    def generate(self, generator: 'BaseGenerator', prompt: str, model_kwargs: Dict[str, Any], num_samples: int = 5, **kwargs) -> str:
+    def generate(self, generator: 'BaseGenerator', prompt: str, model_kwargs: Dict[str, Any], **kwargs, num_samples: int = 5) -> str:
             scored_outputs = []
             for _ in range(num_samples):
-                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(generator.device)).logits.mean().item()
+                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()
                 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(generator.device)
-            outputs = generator.generate(
+    def generate(self, generator: 'BaseGenerator', prompt: str, model_kwargs: Dict[str, Any], **kwargs, num_samples: int = 5) -> str:
+            input_ids = generator.tokenizers["llama"](prompt, return_tensors="pt").input_ids.to(generator.device)
+            outputs = generator.models["llama"].generate(
                 input_ids,
                 num_beams=num_samples,
                 num_return_sequences=num_samples,
                 **model_kwargs
             )
-            return [generator.tokenizer.decode(output, skip_special_tokens=True) for output in outputs]
+            return [generator.tokenizers["llama"].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:
+    def generate(self, generator: 'BaseGenerator', prompt: str, model_kwargs: Dict[str, Any], **kwargs, num_samples: int = 5) -> str:
             results = []
             for _ in range(breadth):
-                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(generator.device)).logits.mean().item()
+                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()
                 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(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(generator.device)).logits.mean().item()
+                    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()
                     results.append((extended_response, score))
             return max(results, key=lambda x: x[1])[0]
         
 
 class COT(GenerationStrategy):
-    def generate(self, generator: 'BaseGenerator', prompt: str, model_kwargs: Dict[str, Any], num_samples: int = 5, **kwargs) -> str:
+    def generate(self, generator: 'BaseGenerator', prompt: str, model_kwargs: Dict[str, Any], **kwargs, num_samples: int = 5) -> str:
         #TODO implement the chain of thought strategy
         
         return "Not implemented yet"       
 
 
 class ReAct(GenerationStrategy):
-    def generate(self, generator: 'BaseGenerator', prompt: str, model_kwargs: Dict[str, Any], num_samples: int = 5, **kwargs) -> str:
+    def generate(self, generator: 'BaseGenerator', prompt: str, model_kwargs: Dict[str, Any], **kwargs, num_samples: int = 5) -> str:
         #TODO implement the ReAct framework       
         return "Not implemented yet"
 #  Add other strategy implementations...