ammarnasr
/

codegen2-1B-security

Text Generation

Model card Files Files and versions

ammarnasr commited on Dec 23, 2023

Commit

5f48ffc

·

1 Parent(s): bbec848

Update handler.py

Files changed (1) hide show

handler.py +1 -1

handler.py CHANGED Viewed

@@ -21,7 +21,7 @@ class EndpointHandler:
           config = {'max_new_tokens':100}
         input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids.to(self.device)
         generated_ids = self.model.generate(input_ids, **config)
-        return self.tokenizer.decode(generated_ids[0], skip_special_tokens=True)

           config = {'max_new_tokens':100}
         input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids.to(self.device)
         generated_ids = self.model.generate(input_ids, **config)
+        return [self.tokenizer.decode(generated_ids[0], skip_special_tokens=True)]