Update modeling_borealis.py

modeling_borealis.py

@@ -145,11 +145,6 @@ class BorealisForConditionalGeneration(PreTrainedModel, PyTorchModelHubMixin):
         )
         return out.loss, out.logits
 
-    def extract_assistant_content(self, text: str) -> str:
-        if "assistant\n" in text:
-            return text.split("assistant\n")[-1].strip()
-        return text.strip()
-
     @torch.no_grad()
     def generate(
         self,
@@ -235,9 +230,9 @@ class BorealisForConditionalGeneration(PreTrainedModel, PyTorchModelHubMixin):
         else:
             txt = self.tokenizer.batch_decode(gen_ids, skip_special_tokens=True)
             if single:
-                return self.extract_assistant_content(txt[0])
+                return txt[0]
             else:
-                return [self.extract_assistant_content(t) for t in txt]
+                return [t for t in txt]
 
     def save_pretrained(self, save_directory, **kwargs):
         os.makedirs(save_directory, exist_ok=True)