Reorder 2 lines

retrieval/pred.py

@@ -64,9 +64,9 @@ def get_pred(model, tokenizer, data, max_length, max_gen, prompt_format, dataset
                 prompt = tokenizer.decode(tokenized_prompt[:half], skip_special_tokens=True)+tokenizer.decode(tokenized_prompt[-half:], skip_special_tokens=True)
             if dataset not in ["trec", "triviaqa", "samsum", "lsht", "lcc", "repobench-p"]: # chat models are better off without build prompt on these tasks
                 prompt = build_chat(tokenizer, prompt, model_name)
-            context_length = input.input_ids.shape[-1]
-            
             input = tokenizer(prompt, truncation=False, return_tensors="pt").to(device)
+            
+            context_length = input.input_ids.shape[-1]
             if dataset == "samsum": # prevent illegal output on samsum (model endlessly repeat "\nDialogue"), might be a prompting issue
                 output = model.generate(
                     **input,