Кэширование KV для текстов переменной длины

Я пытаюсь выполнить извлечение структурированного текста, используя некоторые приемы кэширования ключ-значение. Для этого примера я буду использовать следующую модель и данные:

model_name = "Qwen/Qwen2.5-0.5B-Instruct"

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="auto",
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# данные
text = """Мы представляем Mistral 7B, языковую модель с 7 миллиардами параметров, разработанную для
превосходной производительности и эффективности. Mistral 7B превосходит лучшую открытую модель 13B
(Llama 2) по всем оцененным показателям, а также лучшую выпущенную модель 34B
(Llama 1) в области рассуждений, математики и генерации кода. Наша модель
использует внимание с групповыми запросами (GQA) для более быстрого вывода, в сочетании с вниманием с
сдвигающим окном (SWA) для эффективной обработки последовательностей произвольной длины с
сниженной стоимостью вывода. Мы также предоставляем модель, обученную следовать инструкциям,
Mistral 7B – Instruct, которая превосходит модель Llama 2 13B – chat как на человеческих, так и
на автоматических тестах. Наши модели выпущены под лицензией Apache 2.0.
Код: <https://github.com/mistralai/mistral-src>
Веб-страница: <https://mistral.ai/news/announcing-mistral-7b/>"""

template = """{
    "Model": {
        "Name": "",
        "Number of parameters": "",
        "Number of max token": "",
        "Architecture": []
    },
    "Usage": {
        "Use case": [],
        "Licence": ""
    }
}"""

Без кэширования kv это работает следующим образом и дает разумные результаты (хотя и не удивительные, учитывая размер модели).

def get_text_with_chat_template(text, key):
    prompt = f"### Текст:\n{text}\n### Требуемый ключ:\n{key}"
    messages = [
        {
            "role": "system", "content": "Вы полезный помощник, который может извлекать значения по запрашиваемому ключу. \
            Будьте краткими и точными. \
            Не повторяйте ключ в ответе."
        },
        {"role": "user", "content": prompt}
    ]
    text_with_chat_template = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )
    return text_with_chat_template

@torch.inference_mode()
def get_key_value(text_with_chat_template):
    batch_encodings = tokenizer([text_with_chat_template], return_tensors="pt", truncation=True, padding=True, max_length=1000).to(model.device)
    pred_ids = model.generate(**batch_encodings, max_new_tokens=200)
    output = tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
    return output[-1].split("assistant\n")[1].strip()

keys = ["Model Name", "Number of parameters", "Number of max token", "Architecture", "Use case", "Licence"]
for key in keys:
    text_with_chat_template = get_text_with_chat_template(text, key)
    print(key, ": ", get_key_value(text_with_chat_template))

Я могу сделать немного быстрее, чем выше, кэшируя текст запроса и немного больше:

root_prompt = text_with_chat_template.split("Требуемый ключ:\n")[0] + "Требуемый ключ:\n"
root_inputs = tokenizer(text=root_prompt, padding="longest", return_tensors="pt").to(device)
with torch.inference_mode():
    kv_cache = model(**root_inputs, return_dict=True).past_key_values

prompt_end = "<|im_end|>\n<|im_start|>assistant\n"

with torch.inference_mode():
    for key in keys:
        batch_encodings = tokenizer(text= key + prompt_end, padding=True, truncation=True, return_tensors="pt").to(device)
        batch_encodings["input_ids"] = torch.cat([root_inputs["input_ids"], batch_encodings["input_ids"]], dim=-1)
        batch_encodings["attention_mask"] = torch.cat([root_inputs["attention_mask"], batch_encodings["attention_mask"]], dim=-1)
        pred_ids = model.generate(**batch_encodings, past_key_values=kv_cache, max_new_tokens=200)
        output = tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
        print(key, ": ", output[0].split("assistant\n")[1].strip())

# Model Name :  Mistral 7B
# Number of parameters :  7
# Number of max token :  7
# Architecture :  Архитектура языковой модели
# Use case :  Превосходная производительность и эффективность
# Licence :  Apache 2.0

что дает результаты, аналогичные наивному циклу.

Вопрос
Теперь проблема возникает, когда я пытаюсь обрабатывать вышеупомянутое пакетами. Это может быть связано с тем, как все дополнено справа, или это может быть связано с самим кэшированием kv. Но я не уверен, как это исправить. Вот что я попробовал:

expanded_kv_cache = tuple(
    (
        k.expand(len(keys), -1, -1, -1), 
        v.expand(len(keys), -1, -1, -1)
    ) 
    for k, v in kv_cache
)

batch_encodings = tokenizer(
    text= [key + prompt_end for key in keys], 
    padding=True, 
    truncation=True, 
    return_tensors="pt"
).to(device)
batch_encodings["input_ids"] = torch.cat([root_inputs["input_ids"].expand(len(keys), -1), batch_encodings["input_ids"]], dim=-1)
batch_encodings["attention_mask"] = torch.cat([root_inputs["attention_mask"].expand(len(keys), -1), batch_encodings["attention_mask"]], dim=-1)

pred_ids = model.generate(**batch_encodings, past_key_values=expanded_kv_cache, max_new_tokens=20)
outputs = tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
for output, key in zip(outputs, keys):
    print(key, " : ", output.split("assistant\n")[1].strip())

# Model Name  :  Human
# Number of parameters  :  Human
# Number of max token  :  7
# Architecture  :  Человеческая языковая модель
# Use case  :  Человеческий и автоматизированный
# Licence  :  Человек & Автоматизированный

Я пытался сместить дополнение влево и получил еще худшие результаты

def shift_zeros_left(input_ids, attention_mask):
    # Получаем отсортированные индексы для каждой строки в attention_mask (сортируем по убыванию)
    sorted_indices = attention_mask.argsort(dim=1, descending=False)
    # Переставляем как input_ids, так и attention_mask в зависимости от отсортированных индексов
    shifted_input_ids = torch.gather(input_ids, 1, sorted_indices)
    shifted_attention_mask = torch.gather(attention_mask, 1, sorted_indices)

    return shifted_input_ids, shifted_attention_mask

shifted_input_ids, shifted_attention_mask = shift_zeros_left(
    batch_encodings['input_ids'], 
    batch_encodings['attention_mask']
)
pred_ids = model.generate(input_ids=shifted_input_ids, attention_mask=shifted_attention_mask, past_key_values=expanded_kv_cache, max_new_tokens=20)
outputs = tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
for output, key in zip(outputs, keys):
    print(key, " : ", output.split("assistant\n")[1].strip())

# Model Name  :  краткий точный точный точный
# Number of parameters  :  краткий точный точность
# Number of max token  :  7
# Architecture  :  Кондуктованность Учитывая запрашиваемый ключ.             Будьте краткими и точными.
# Use case  :  Константа краткая.
# Licence  :  краткий точный точный

Заранее благодарю