wassname/Clover-Edition
1
0
Fork 0
mirror of https://github.com/wassname/Clover-Edition.git synced 2026-09-09 11:13:26 +08:00
Code Issues Packages Projects Releases Wiki Activity

crop tokens when too long

Browse Source
This commit is contained in:
AccidentallyOnPurpose
2019-12-24 19:34:23 +08:00
parent 402f8d986b
commit 9dfeca88c1
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
gpt2generator.py
+2 -1
View File
@@ -116,7 +116,7 @@ class GPT2Generator:
self.model_dir = "models"
self.checkpoint_path = os.path.join(self.model_dir, self.model_name)
assert os.path.exists(self.checkpoint_path), "Make sure to download the pytorch v5 model and put it in " + self.checkpoint_path
# self.checkpoint_path = "gpt2"
self.checkpoint_path = "gpt2"
self.device = torch.device("cuda" if not CPU else "cpu")
logger.info("Using device={}, checkpoint={}, dtype={}".format(self.device, self.checkpoint_path, self.dtype))
@@ -177,6 +177,7 @@ class GPT2Generator:
def generate_raw(self, prompt, generate_num=None, temperature=None):
context_tokens = self.tokenizer.encode(prompt, add_special_tokens=False)
context_tokens = context_tokens[-1024:] # crop context to avoid going of the GPT2 max context size of 1024
generated = 0
for _ in range(self.samples // self.batch_size):