Update services/strategy.py

services/strategy.py

@@ -31,7 +31,7 @@ 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.model.generate(input_ids, **model_kwargs)
+        output = generator.generate(input_ids, **model_kwargs)
         return generator.tokenizer.decode(output[0], skip_special_tokens=True)
 
 
@@ -40,7 +40,7 @@ class MajorityVotingStrategy(GenerationStrategy):
         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)
+            output = generator.generate(input_ids, **model_kwargs)
             outputs.append(generator.tokenizer.decode(output[0], skip_special_tokens=True))
         return max(set(outputs), key=outputs.count)
 
@@ -66,7 +66,7 @@ class BeamSearch(GenerationStrategy):
                 num_return_sequences=num_samples,
                 **model_kwargs
             )
-            return [self.llama_tokenizer.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):