ammarnasr
/

codegen2-1B-security

Text Generation

Model card Files Files and versions

ammarnasr commited on Dec 23, 2023

Commit

6f0551a

·

1 Parent(s): 5677d77

Update handler.py

Files changed (1) hide show

handler.py +1 -1

handler.py CHANGED Viewed

@@ -17,7 +17,7 @@ class EndpointHandler:
         prompt = data["inputs"]
         input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids.to(self.device)
         generated_ids = self.model.generate(input_ids)
-        return self.tokenizer.decode(generated_ids[0], skip_special_tokens=True)

         prompt = data["inputs"]
         input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids.to(self.device)
         generated_ids = self.model.generate(input_ids)
+        return [self.tokenizer.decode(generated_ids[0], skip_special_tokens=True)]