better handling or eos tokens, hard limit on vRAM usage (for RX 7900XTX)

FontaineRiant · Mar 21, 2024 · 3d2b0e5 · 3d2b0e5
1 parent 5bc58ee
commit 3d2b0e5
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/generator/generator.py b/generator/generator.py
@@ -28,7 +28,7 @@ def __init__(self,
         self.enc = AutoTokenizer.from_pretrained(model_name, add_prefix_space=False)
 
         self.length = length
-        self.max_history = self.model.config.max_position_embeddings - self.length
+        self.max_history = min(self.model.config.max_position_embeddings - self.length, 6000)
 
     def __del__(self):
         pass

diff --git a/story/story.py b/story/story.py
@@ -74,8 +74,8 @@ def clean_input(self, action=''):
         return text.strip()
 
     def clean_result(self, result):
-        result = re.sub(rf'^({self.gen.model.config.prefix})+', '', result) # remove leading endoftext tokens
-        result = re.sub(rf'{self.gen.model.config.prefix}[\s\S]*$', '', result)  # parse endoftext token that end the text
+        result = re.sub(rf'^({self.gen.enc.eos_token})+', '', result) # remove leading endoftext tokens
+        result = re.sub(rf'{self.gen.enc.eos_token}[\s\S]*$', '', result)  # parse endoftext token that end the text
 
         result = result.replace("’", "'")
         result = result.replace("`", "'")