Attempt fix for RuntimeError: p.attn_bias_ptr is not correctly aligned

main.py

@@ -1,29 +1,28 @@
-from langchain_core.messages import SystemMessage
-from langchain_core.prompts import ChatPromptTemplate
-from langchain_core.runnables import RunnablePassthrough
+from transformers import AutoModelForCausalLM, AutoTokenizer
 from langchain_core.output_parsers import StrOutputParser
+from langchain_core.runnables import RunnablePassthrough
+from langchain_core.prompts import ChatPromptTemplate
 from langgraph.checkpoint.memory import MemorySaver
+from langchain_core.messages import SystemMessage
 from langchain_core.messages import ToolMessage
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-from datetime import datetime
 from dotenv import load_dotenv
-import os
-import glob
-import torch
+from datetime import datetime
 import logging
+import torch
+import glob
 import ast
+import os
 
-# To use OpenAI models (remote)
-from langchain_openai import ChatOpenAI
-
-# To use Hugging Face models (local)
+# Imports for local and remote chat models
 from langchain_huggingface import ChatHuggingFace, HuggingFacePipeline
+from langchain_openai import ChatOpenAI
 
 # Local modules
-from index import ProcessFile
+from pipeline import MyTextGenerationPipeline
 from retriever import BuildRetriever, db_dir
-from graph import BuildGraph
 from prompts import answer_prompt
+from index import ProcessFile
+from graph import BuildGraph
 
 # -----------
 # R-help-chat
@@ -157,16 +156,17 @@ def GetChatModel(compute_mode, ckpt_dir=None):
             torch_dtype=torch.bfloat16,
         )
 
-        # ToolCallingLLM needs return_full_text=False in order to parse just the assistant response;
-        # the JSON function descriptions in the full response cause an error in ToolCallingLLM
-        pipe = pipeline(
-            "text-generation",
+        # Use MyTextGenerationPipeline with custom preprocess() method
+        pipe = MyTextGenerationPipeline(
             model=model,
             tokenizer=tokenizer,
+            # ToolCallingLLM needs return_full_text=False in order to parse just the assistant response
             return_full_text=False,
             # It seems that max_new_tokens has to be specified here, not in .invoke()
-            max_new_tokens=1000,
+            max_new_tokens=2000,
         )
+        # We need the task so HuggingFacePipeline can deal with our class
+        pipe.task = "text-generation"
 
         llm = HuggingFacePipeline(pipeline=pipe)
         chat_model = ChatHuggingFace(llm=llm)

mods/bm25s_retriever.py

@@ -155,6 +155,7 @@ class BM25SRetriever(BaseRetriever):
         *,
         run_manager: CallbackManagerForRetrieverRun,
     ) -> List[Document]:
+        # from bm25s import tokenize as bm25s_tokenize
         from mods.bm25s_tokenization import tokenize as bm25s_tokenize
 
         processed_query = bm25s_tokenize(query, return_ids=False)

pipeline.py

@@ -0,0 +1,86 @@
+from transformers.pipelines.text_generation import Chat
+from transformers import TextGenerationPipeline
+from typing import Dict
+
+
+class MyTextGenerationPipeline(TextGenerationPipeline):
+    """
+    This subclass overrides the preprocess method to add pad_to_multiple_of=8 to tokenizer_kwargs.
+    Fix for: "RuntimeError: p.attn_bias_ptr is not correctly aligned"
+    https://github.com/google-deepmind/gemma/issues/169
+    """
+
+    def preprocess(
+        self,
+        prompt_text,
+        prefix="",
+        handle_long_generation=None,
+        add_special_tokens=None,
+        truncation=None,
+        padding=None,
+        max_length=None,
+        continue_final_message=None,
+        **generate_kwargs,
+    ):
+        print(f"PADDING: {padding}")
+        # Only set non-None tokenizer kwargs, so as to rely on the tokenizer's defaults
+        tokenizer_kwargs = {
+            "add_special_tokens": add_special_tokens,
+            "truncation": truncation,
+            "padding": padding,
+            "max_length": max_length,
+            "pad_to_multiple_of": 8,
+        }
+        tokenizer_kwargs = {
+            key: value for key, value in tokenizer_kwargs.items() if value is not None
+        }
+
+        if isinstance(prompt_text, Chat):
+            tokenizer_kwargs.pop(
+                "add_special_tokens", None
+            )  # ignore add_special_tokens on chats
+            # If the user passes a chat that ends in an assistant message, we treat it as a prefill by default
+            # because very few models support multiple separate, consecutive assistant messages
+            if continue_final_message is None:
+                continue_final_message = prompt_text.messages[-1]["role"] == "assistant"
+            inputs = self.tokenizer.apply_chat_template(
+                prompt_text.messages,
+                add_generation_prompt=not continue_final_message,
+                continue_final_message=continue_final_message,
+                return_dict=True,
+                return_tensors=self.framework,
+                **tokenizer_kwargs,
+            )
+        else:
+            inputs = self.tokenizer(
+                prefix + prompt_text, return_tensors=self.framework, **tokenizer_kwargs
+            )
+
+        inputs["prompt_text"] = prompt_text
+
+        if handle_long_generation == "hole":
+            cur_len = inputs["input_ids"].shape[-1]
+            if "max_new_tokens" in generate_kwargs:
+                new_tokens = generate_kwargs["max_new_tokens"]
+            else:
+                new_tokens = (
+                    generate_kwargs.get("max_length", self.generation_config.max_length)
+                    - cur_len
+                )
+                if new_tokens < 0:
+                    raise ValueError("We cannot infer how many new tokens are expected")
+            if cur_len + new_tokens > self.tokenizer.model_max_length:
+                keep_length = self.tokenizer.model_max_length - new_tokens
+                if keep_length <= 0:
+                    raise ValueError(
+                        "We cannot use `hole` to handle this generation the number of desired tokens exceeds the"
+                        " models max length"
+                    )
+
+                inputs["input_ids"] = inputs["input_ids"][:, -keep_length:]
+                if "attention_mask" in inputs:
+                    inputs["attention_mask"] = inputs["attention_mask"][
+                        :, -keep_length:
+                    ]
+
+        return inputs