Update services/strategy.py

services/strategy.py

@@ -30,13 +30,17 @@ class GenerationStrategy(ABC):
 class DefaultStrategy(GenerationStrategy):
     @observe()
     def generate(self, generator: 'BaseGenerator', prompt: str, model_kwargs: Dict[str, Any], **kwargs) -> str:
-        
-        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)
+        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)
+    #def generate(self, generator: 'BaseGenerator', prompt: str, model_kwargs: Dict[str, Any], **kwargs) -> str:
+    #    
+    #    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):