From 48b3f71ed49d7f63140ada0e2281c4fc819c1367 Mon Sep 17 00:00:00 2001
From: cornzz <39997278+cornzz@users.noreply.github.com>
Date: Thu, 14 Nov 2024 19:57:00 +0100
Subject: [PATCH] Fix(LLMLingua): fix prompts smaller than iterative_size not
 being compressed (#196)

---
 llmlingua/prompt_compressor.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/llmlingua/prompt_compressor.py b/llmlingua/prompt_compressor.py
index da5e765..b9a0ea9 100644
--- a/llmlingua/prompt_compressor.py
+++ b/llmlingua/prompt_compressor.py
@@ -1413,6 +1413,8 @@ def get_compressed_input(
         self_input_ids=None,
         self_attention_mask=None,
     ):
+        if end < iterative_size:
+            end = iterative_size
         if self_loss is not None:
             need_idx = torch.concat(
                 [