Skip decoding special tokens (#1244)

SilasMarvin · web-flow · commit da509e28602b · 2023-12-18T14:53:14.000-08:00
diff --git a/pgml-extension/src/bindings/transformers/transformers.py b/pgml-extension/src/bindings/transformers/transformers.py
@@ -314,6 +314,7 @@ def stream(self, input, timeout=None, **kwargs):
                 self.tokenizer,
                 timeout=timeout,
                 skip_prompt=True,
+                skip_special_tokens=True
             )
             if "chat_template" in kwargs:
                 input = self.tokenizer.apply_chat_template(
@@ -338,6 +339,7 @@ def stream(self, input, timeout=None, **kwargs):
             streamer = TextIteratorStreamer(
                 self.tokenizer,
                 timeout=timeout,
+                skip_special_tokens=True
             )
             input = self.tokenizer(input, return_tensors="pt", padding=True).to(
                 self.model.device
