diff --git "a/query_rewrite/README.md" "b/query_rewrite/README.md" new file mode 100644--- /dev/null +++ "b/query_rewrite/README.md" @@ -0,0 +1,287 @@ +--- +license: apache-2.0 +language: +- en +pipeline_tag: text-generation +library_name: peft +library_name: transformers +--- +--- + +# Query Rewrite Intrinsic + +## Model Summary + + + +We are releasing a query rewrite intrinsic, implemented as a family of adapters that are fine-tuned specifically for the following task: + + Given a multi-turn conversation between a user and an AI assistant, rewrite the last + user utterance (query) by transforming it (only if necessary) into an equivalent version that + is standalone and can be understood by itself (without the conversation). + +While the intrinsic is general purpose, one of its main use cases is in RAG settings where the ability to rewrite a user query into a standalone version directly improves the retriever performance, which in turn improves the answer generation performance. Outside of RAG, there are other conversational use cases that require to rewrite a user query, for example, before accessing a database, or before routing to other APIs or tools. The intrinsic does not need any RAG documents (which may be present in the context, in a RAG setting) and uses only the dialog turns with what is being said between the user and assistant. We are providing experimental results in a RAG setting as well as on a specialized enterprise (non-RAG) setup, showing in both cases that the intrinsic performance is significantly higher than when prompting out-of-the-box models, including open-source models such as gpt-oss as well as frontier models such as gpt-4o. + +The adapters released here work with both the IBM granite-3.3 family of models (2b and 8b) as well as gpt-oss-20b. + +- **Developer:** IBM Research +- **Model type:** LoRA adapter for [ibm-granite/granite-3.3-2b-instruct](https://huggingface.co/ibm-granite/granite-3.3-2b-instruct) and [ibm-granite/granite-3.3-8b-instruct](https://huggingface.co/ibm-granite/granite-3.3-8b-instruct) +- **License:** [Apache 2.0](https://www.apache.org/licenses/LICENSE-2.0) + + +## Intended use + +The intrinsic gives the ability to rewrite the last user query in a multi-turn conversation. Typically, the rewrite is a form of expansion that inlines into the query any implicit references that are made to entities, concepts, or even parts of the conversation that occur in the previous turns (either by the user or the AI assistant). Such expansion can include coreference resolution (i.e., replacement of pronouns with the actual entities), handling of ellipsis, which is the common linguistic phenomenon where parts of a sentence or phrase are omitted by the user, but can be understood from the context (i.e., for whom, of what, with respect to something discussed above, etc.). + +As a result of the expansion, the query becomes a standalone query, still equivalent in meaning with what the user asked in the last turn. The rewritten query can be sent to downstream tasks (e.g., to a retriever in a RAG setting) as a better replacement for the original user query, and without the need for (a potentially very long) context. + +> [!TIP] +> Note: While you can invoke the query rewrite intrinsic directly, it is strongly recommended to call it through [granite-common](https://github.com/ibm-granite/granite-common), which wraps the model with a tailored I/O processor, enabling a friendlier development interface. The I/O processor takes care of several data transformation and validation tasks that would be otherwise required (including sending the appropriate instruction for the intrinsic) as well as validating the intrinsic's output. We next describe the input/output of the query rewrite intrinsic when invoked through granite-common. + +**Intrinsic input**: The input to the query rewrite intrinsic is an OpenAI-compatible chat completion request, containing a list of conversation turns that can alternate between the `user` and `assistant` role, and ending with a `user` turn. The last `user` turn in the list is assumed to be the query that needs to be rewritten (if not already standalone). + +**Intrinsic output**: The output of the query rewrite intrinsic is the result of the original chat completion request formatted as a +JSON object as follows: +```json +{ + "rewritten_question": +} +``` + +Please see the code snippets in the Quickstart Example section below for examples that illustrate the intrinsic's input/output. + +## Quickstart Example Using Granite-Common + +To run the intrinsic using an OpenAI-compatible inference backend, such as vLLM, follow the steps below. We recommend using Python 3.11 or higher. + +1. Install the granite-common library: + ``` + pip install granite-common[nltk] + ``` + +2. Install the Hugging Face CLI: + ``` + pip install -U "huggingface_hub[cli]" + ``` + +3. Install vLLM: + ``` + pip install vllm + ``` + +4. Download the intrinsics library: + ``` + hf download ibm-granite/rag-intrinsics-lib --local-dir ./rag-intrinsics-lib + cd rag-intrinsics-lib + ``` + +5. Edit if needed the vLLM startup script found in `./rag-intrinsics-lib/run_vllm.sh` using your favorite editor: + + Edit the constants `BASE_MODEL_NAME` and `BASE_MODEL_ORG` depending on the base model on which the desired adapter has been trained. Optionally, edit the constant `PORT` to change the port on which vLLM will run. Save the modified file and exit the editor. + +6. [Optional] If you are on a cluster, allocate a node with GPU, which will be used as inference node. On a CPU, things will be slower. Edit the below LSF command as needed (e.g., among other things, you will need to change `-q gpu` to reflect the actual GPU queue that you have access to). Also, before that, you will need to make sure you activate the same python environment where you installed granite-common and vllm in the steps above. + ``` + bsub -q gpu -gpu "num=1" -hl -n 1 -W 12:0 -Is bash + hostname -f # get the node name + ``` + +7. Start vLLM through the startup script. If you are on a cluster and have allocated the GPU node, do that from the shell of the GPU node (e.g., started with the `bsub` job above). The first time you run the script, you may have to change the permissions to allow execution: + ``` + chmod u+x ./run_vllm.sh + ./run_vllm.sh & + ``` + +8. Run the following Python code snippet. This does not have to be on the same node as the one you are running the vLLM server and, in particular, it can be on a regular CPU node. + + ``` + import json + import openai + import granite_common + intrinsic_name = "query_rewrite" + # Change the following constant as needed to select a different base model + base_model_name = "granite-3.3-8b-instruct" + # Change the following constants as needed to reflect the location of the vLLM server. + # For example, if you are on a cluster and you have allocated the GPU node as above, use + # the nodename obtained above (with `hostname -f`) instead of localhost. + # The selected port should be identical to the one you specified in the vLLM startup script + openai_base_url = "http://localhost:55555/v1" + openai_api_key = "rag_intrinsics_1234" + # Fetch IO configuration file from Hugging Face Hub + io_yaml_file = granite_common.intrinsics.util.obtain_io_yaml( + intrinsic_name, base_model_name + ) + # Instantiate input/output processors + rewriter = granite_common.IntrinsicsRewriter(config_file=io_yaml_file) + result_processor = granite_common.IntrinsicsResultProcessor(config_file=io_yaml_file) + # Sample request + request_json = { + "messages": [ + { + "role": "assistant", + "content": "Welcome to pet questions!" + }, + { + "role": "user", + "content": "I have two pets, a dog named Rex and a cat named Lucy." + }, + { + "role": "assistant", + "content": "Great, what would you like to share about them?" + }, + { + "role": "user", + "content": "Rex spends a lot of time in the backyard and outdoors, and Lucy is always inside." + }, + { + "role": "assistant", + "content": "Sounds good! Rex must love exploring outside, while Lucy probably enjoys her cozy indoor life." + }, + { + "role": "user", + "content": "But is he more likely to get fleas because of that?" + } + ] + } + # Add other parameters + request_json["model"] = intrinsic_name + request_json["temperature"] = 0.0 + # Apply input processor + rewritten_request = rewriter.transform(request_json) + # Run inference + client = openai.OpenAI(base_url=openai_base_url, api_key=openai_api_key) + chat_completion = client.chat.completions.create(**rewritten_request.model_dump()) + # Apply output processor + processed_chat_completion = result_processor.transform( + chat_completion, rewritten_request + ) + # Verify that the contents of the completion is valid JSON and pretty-print the JSON. + parsed_contents = json.loads(processed_chat_completion.choices[0].message.content) + print("JSON output:") + print(json.dumps(parsed_contents, indent=2)) + ``` + +The post-processed JSON output for the above example should be something similar to this: +```json +{ + "rewritten_question": "Is Rex, my outdoor dog, more likely to get fleas because of his time spent in the backyard and outdoors?" +} +``` + +## Training Details + +The training data contains both: 1) standalone examples, which teach the intrinsic to refrain from rewriting user questions that are already standalone, and 2) non-standalone examples containing a diversity of patterns that are used to teach the intrinsic to expand the user turn so that it becomes standalone. + +### Training Data + +The training data uses the publicly available Cloud corpus of technical documentation pages from [MT-RAG](https://arxiv.org/abs/2501.03468). Based on this corpus of documents, we constructed a dataset consisting of high-quality, human-created conversations, where the last turn of the conversation comes into versions: non-standalone version, and corresponding standalone version. In addition, we have also used a synthetically generated set of training examples, to maximize the diversity across a variety of patterns. + +The training dataset is proprietary and was obtained in combination with a third-party company who contracted the human annotators. + +### Robustness to System Prompts + +In a typical Retrieval-Augmented Generation (RAG) setup, different researchers or practitioners may use various system prompts tailored to their specific use cases. For the IBM granite versions of the adapter, to enhance robustness against these variations, we generated three distinct versions of each training sample, each paired with a different system prompt. This expanded and diversified training dataset is then used to train the LoRA adapters, improving their ability to handle diverse prompt styles effectively. + +System Prompts Used (specifically for IBM granite base models): + +Version 1: <|start_of_role|>system<|end_of_role|> Knowledge Cutoff Date: April 2024. Today's Date: May 20, 2025. You are Granite, developed by IBM. You are a helpful AI assistant. <|end_of_text|> + +Version 2: <|start_of_role|>system<|end_of_role|> Knowledge Cutoff Date: April 2024. Today's Date: May 20, 2025. You are Granite, developed by IBM. Write the response to the user's input by strictly aligning with the facts in the provided documents. If the information needed to answer the question is not available in the documents, inform the user that the question cannot be answered based on the available data.<|end_of_text|> + +Version 3: An empty system prompt (no instructions provided). + +This approach ensures that the adapters remain effective and reliable across varying system prompt formats commonly encountered in real-world RAG applications. + +#### Training Hyperparameters + +The adapters, both for granite and gpt-oss, were fine-tuned using PEFT under the following regime: rank = 32, learning rate = 3e-6, number of epochs = 25, and linear learning rate scheduler. + + +## Evaluation +We provide two types of evaluation: 1) evaluation that is specific for the task of query rewrite, 2) end-to-end evaluation in a RAG setting, where we evaluate how query rewrite impacts: a) the retriever performance, and b) the final answer generation. + +### Evaluation Specific for the Query Rewrite Task + +Here we evaluate the quality of the rewritten queries themselves on an enterprise internal dataset. This is a dataset with two turn conversations, where the last user turn may or may not be standalone. We do have the gold rewritten queries, and also make use of a LLM judge (with LLama-3.3-70b as the model) with a specific prompt to check whether the model-generate rewriting is equivalent to the gold rewriting. This is a challenging benchmark, with the specific requirement that the models minimally change the query with only the additions needed to make the query standalone (it not already standalone). If any changes are not minimal, the judge will penalize the rewritten query. + +alt text + +In the chart above, towards the left, we show the performance of the LoRA adapters for three base models: gpt-oss-20b, granite-3.3-8b-instruct, and granite-3.3-2b-instruct. We show how these compare to the versions where we just prompt the same three base models, using just the query rewrite instruction. Towards the right, we show results when we prompt two frontier models (gpt-4o, and gpt-4o-mini) as well as the larger gpt-oss-120b, using just the query rewrite instruction. We also include, in the middle, the performance of an out-of-the-box granite base model when prompted with a custom prompt that includes 13 examples of query rewrites, in addition to the query rewrite instruction. We see that all of these are dominated by the fine-tuned LoRA adapters (even by the smaller 2b version). Furthermore, we see that the performance of the adapters is quite stable across the two families of models (gpt-oss and ibm-granite) and across the different model sizes. + + +### Evaluation of Retriever + +We evaluate Recall@k on the [MT-RAG](https://arxiv.org/abs/2501.03468) benchmark, under various query rewrite strategies for the retriever. All retrieved passages are obtained using the Elser retriever with the same settings as in the above paper. In addition to the LoRA adapter, we include several other baselines, including no-rewrite (where we send the last user turn to the retriever as-is), Mixtral rewrites, as well as gold rewrites (human-created). + +We evaluate on three different testsets: a) full MT-RAG dataset (842 data points with last user turns); b) the non-standalone subset of MT-RAG dataset, which is a subset of 260 (out of 842) last user turns that were annotated by humans as non-standalone (i.e., they are dependent on the prior context); c) the standalone subset of MT-RAG dataset, which is the complementary subset, with all the last user turns that were annotated by humans as standalone. + +a. Evaluation of Recall@k on full MT-RAG dataset. + +| Strategy | Recall@5 | Recall@10 | Recall@20 | +| --------------------------- | --------------- | ----------------- | ------------ | +| No rewrite | 0.486 | 0.587 | 0.665 | +| Mixtral 8x7b rewrite | 0.522 | 0.642 | 0.72 | +| Gold rewrite | 0.563 | 0.674 | 0.747 | +| Granite 3.3-8b LoRA rewrite | 0.563 | 0.682 | 0.762 | + +b. Evaluation of Recall@k on the non-standalone subset of MT-RAG. + +| Strategy | Recall@5 | Recall@10 | Recall@20 | +| --------------------------- | --------------- | ----------------- | ------------ | +| No rewrite | 0.263 | 0.338 | 0.435 | +| Mixtral 8x7b rewrite | 0.362 | 0.488 | 0.574 | +| Gold rewrite | 0.479 | 0.582 | 0.662 | +| Granite 3.3-8b LoRA rewrite | 0.445 | 0.556 | 0.648 | + +c. Evaluation of Recall@k on the standalone subset of MT-RAG. + +| Strategy | Recall@5 | Recall@10 | Recall@20 | +| --------------------------- | --------------- | ----------------- | ------------ | +| No rewrite | 0.609 | 0.723 | 0.792 | +| Mixtral 8x7b rewrite | 0.613 | 0.733 | 0.809 | +| Gold rewrite | 0.609 | 0.723 | 0.792 | +| Granite 3.3-8b LoRA rewrite | 0.628 | 0.751 | 0.824 | + +If we focus on Recall@20 numbers, as one instance of the metric, there is an overall 9.7 percentage points jump when using query rewrite with the Granite 3.3-8b LoRA adapter versus when using the no rewrite strategy. This jump is more pronounced on the non-standalone fragment, where query rewrite with the Granite 3.3-8b LoRA adapter leads to almost 21 percentage points improvement over the no-rewrite strategy. Also, we can observe that the numbers with the LoRA rewrites are very close to what can be obtained with the gold rewrites on non-standalones (and slightly better on standalones for LoRA -- human annotators were instructed to leave the query unchanged when classifying it as standalone, however, the LoRA adapter may still perform some rewriting which turns out to further improve the recall). + +### Evaluation of Answer Generation + +We evaluate answer generation quality, with top-k passages retrieved under the various query rewrite strategies for the retriever. We choose here k = 20, but similar trends take place for other values of k. We used Granite-3.3-8b instruct as the answer generator, and [RAGAS](https://arxiv.org/abs/2309.15217) Faithfulness on the answerable subset of MT RAG data, [JAFS](https://arxiv.org/abs/2504.11704) that rewards the model for correctly abstaining on unanswerable queries (full credit) and for +providing faithful answers on answerable queries (partial credit based on RAGAS Faithfulness), and [RAD-Bench](https://arxiv.org/abs/2409.12558) score as metrics for answer quality. We use the same three testsets as above. + +a. Evaluation of answer quality on full MT-RAG dataset. + +| Strategy | RAGAS-F (Answerable Subset) | RAD-Bench | JAFS | +| --------------------------- | ---------------- | ---------------- | ---------------- | +| No rewrite | 0.793 | 0.678 | 0.664 | +| Mixtral 8x7b rewrite | 0.78 | 0.679 | 0.682 | +| Gold rewrite | 0.81 | 0.686 | 0.67 | +| Granite 3.3-8b LoRA rewrite | 0.874 | 0.698 | 0.722 | + +b. Evaluation of answer quality on non-standalone MT-RAG subset. + +| Strategy | RAGAS-F (Answerable Subset) | RAD-Bench | JAFS | +| --------------------------- | ---------------- | ---------------- | ---------------- | +| No rewrite | 0.695 | 0.618 | 0.581 | +| Mixtral 8x7b rewrite | 0.776 | 0.644 | 0.627 | +| Gold rewrite | 0.786 | 0.661 | 0.634 | +| Granite 3.3-8b LoRA rewrite | 0.865 | 0.669 | 0.70 | + +c. Evaluation of answer quality on standalone subset of MT-RAG. + +| Strategy | RAGAS-F (Answerable Subset) | RAD-Bench | JAFS | +| --------------------------- | ---------------- | ---------------- | ---------------- | +| No rewrite | 0.845 | 0.71 | 0.708 | +| Mixtral 8x7b rewrite | 0.854 | 0.697 | 0.71 | +| Gold rewrite | 0.845 | 0.71 | 0.708 | +| Granite 3.3-8b LoRA rewrite | 0.88 | 0.713 | 0.734 | + + + +As with Recall, similar observations can be made here as well. Specifically, on the full dataset, we see an 8.1 percentage points jump in RAGAS Faithfulness (from 0.793 to 0.874), a 2 percentage points jump in RAD-Bench score (from 0.678 to 0.698), and a 5.8 percentage points jump in JAFS (from 0.664 to 0.722) when using query rewrite with the Granite 3.3-8b LoRA adapter versus when using the no rewrite strategy. This improvement is more pronounced on the non-standalone subset, where query rewrite with the Granite 3.3-8b LoRA adapter leads to a 17 percentage points jump in RAGAS Faithfulness (from 0.695 to 0.865), a 5.1 percentage points jump in RAD-Bench score (from 0.618 to 0.669), and an 11.9 percentage points jump in JAFS (from 0.581 to 0.70). + +## Contact +[Lucian Popa](mailto:lpopa@us.ibm.com) + +### Framework versions + +- PEFT 0.17.2