Update true_poetry.py

fix for Huggingface v4
summerstay · Apr 30, 2021 · a726f85 · a726f85
1 parent 83ff4e2
commit a726f85
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/true_poetry.py b/true_poetry.py
@@ -377,7 +377,8 @@ def expand_node(sentence, past):
         input_ids = torch.tensor([sentence[-1]]).unsqueeze(0)
     inputs = {'input_ids': input_ids}    
     with torch.no_grad():
-        logits, past = model(**inputs, past=past)
+        logits, past = model(**inputs, past_key_values=past, return_dict=False)
+        // use this line for Huggingface v3: logits, past = model(**inputs, past=past)
         logits[0][0][50256]=-math.inf # no <end of text> token
         logits = logits[:, -1, :]  
         probs = F.softmax(logits, dim=-1).tolist()[0]
@@ -565,4 +566,4 @@ def poem_scheme(kind):
             poem_line[line] = poem_line[line][1:]
     for line in range(1,number_of_lines):
         print(tokenizer.decode(poem_line[line]))
-
+