QueryRefiner-0.5B-v0.1-GRPO

QueryRefiner-0.5B-v0.1-GRPO

<!-- RAG Query Analyzer 0.5B v0.1 -->

Model Description

QueryRefiner-0.5B-v0.1-GRPO is a specialized query analysis and rephrasing model fine-tuned from Qwen2.5-0.5B-Instruct for Retrieval-Augmented Generation (RAG) systems. The model has been trained to analyze user queries through a structured, ontology-inspired reasoning framework that produces XML outputs with domain classification, intent detection, conceptual analysis, and optimized query reformulations designed to improve document retrieval precision.

This model is part of a two-model series addressing a critical challenge in RAG systems: converting ambiguous, vague, or suboptimally-phrased user queries into retrieval-friendly formats through systematic decomposition and analysis. By employing a structured reasoning approach that explicitly identifies domains, intents, concepts, and relationships, the model produces higher-quality rephrasings informed by comprehensive query understanding. <!-- Rather than directly answering questions, the model performs query understanding and transformation, making it an ideal preprocessing component for information retrieval pipelines. -->

Note: This model is currently in an experimental, beta stage. It represents an initial release focused on establishing baseline capabilities. Future iterations may address current limitations, expand domain coverage, or incorporate additional output features based on real-world deployment feedback.

Examples

<details> <summary><i>Unroll to see Example #1</i></summary>

User query :

How can artificial intelligence be utilized to improve emergency response systems in urban environments?

Model output:

<query_analysis>
  <domain ambiguous="false">
    <candidate confidence="1.0">Technology</candidate>
  </domain>
  <intent ambiguous="false">
    <candidate confidence="1.0">Informational</candidate>
  </intent>
  <concepts>
    <explicit>artificial intelligence, emergency response systems, urban environments</explicit>
    <implicit>safety, efficiency, decision-making, urban planning</implicit>
  </concepts>
  <relations>
    <relation subject="artificial intelligence" predicate="utilized to improve" object="emergency response systems"/>
    <relation subject="emergency response systems" predicate="used in" object="urban environments"/>
  </relations>
  <insufficient_context>false</insufficient_context>
  <rephrased>true</rephrased>
  <rephrased_query>how can artificial intelligence technologies be leveraged to enhance emergency management and response in urban settings</rephrased_query>
</query_analysis>

</details> <details> <summary><i>Unroll to see Example #2</i></summary>

User query:

How have its fiscal policies regarding these things impacted its overall financial stability and future outlook, particularly concerning external market fluctuations?

Model output:

<query_analysis>
  <domain ambiguous="true">
    <candidate confidence="0.6">Finance</candidate>
    <candidate confidence="0.4">Economics</candidate>
  </domain>
  <intent ambiguous="true">
    <candidate confidence="0.7">Impact assessment</candidate>
    <candidate confidence="0.3">Analysis of policy</candidate>
  </intent>
  <concepts>
    <explicit>fiscal policies, external market fluctuations, financial stability, future outlook</explicit>
    <implicit>[ENTITY]</implicit>
  </concepts>
  <relations>
    <relation subject="[ENTITY]" predicate="has" object="fiscal policies"/>
    <relation subject="[ENTITY]" predicate="impacts" object="financial stability"/>
    <relation subject="[ENTITY]" predicate="impacts" object="future outlook"/>
    <relation subject="[ENTITY]" predicate="relates to" object="external market fluctuations"/>
  </relations>
  <insufficient_context>false</insufficient_context>
  <rephrased>true</rephrased>
  <rephrased_query>How have [ENTITY]'s fiscal policies regarding external market fluctuations impacted its overall financial stability and future outlook?</rephrased_query>
</query_analysis>

</details> <details> <summary><i>Unroll to see Example #3</i></summary>

User query :

How does she describe her emotional journey and the complexities of her interactions with him in that book?

Model output:

<query_analysis>
  <domain ambiguous="true">
    <candidate confidence="0.5">Literature</candidate>
    <candidate confidence="0.5">Psychology</candidate>
  </domain>
  <intent ambiguous="true">
    <candidate confidence="0.6">Analysis</candidate>
    <candidate confidence="0.4">Description</candidate>
  </intent>
  <concepts>
    <explicit>emotional journey, interactions, book</explicit>
    <implicit>[PERSON], complex</implicit>
  </concepts>
  <relations>
    <relation subject="[PERSON]" predicate="describe" object="emotional journey"/>
    <relation subject="[PERSON]" predicate="interact with" object="him"/>
    <relation subject="book" predicate="describe" object="[PERSON]'s journey and interactions"/>
  </relations>
  <insufficient_context>true</insufficient_context>
  <rephrased>true</rephrased>
  <rephrased_query>How does [PERSON] describe her emotional journey and the complexities of her interactions with [PERSON] in the book?</rephrased_query>
</query_analysis>

</details>

Intended Use

The model is designed to serve as a query analysis layer in RAG systems, where it processes user inputs before they are sent to retrieval components. It can handle queries with varying degrees of ambiguity and identify when queries lack sufficient context for effective retrieval. The structured XML output enables downstream systems to make informed decisions about how to handle each query based on confidence scores, ambiguity flags, and the rephrased version.

Typical deployment scenarios include conversational search systems, question-answering platforms, document retrieval services, and any application where understanding user intent and optimizing queries for semantic search is critical. The model's compact 0.5B parameter size makes it suitable for deployment in resource-constrained environments or as part of larger multi-component systems where latency is a concern.

Training Data

The model was trained on two complementary datasets, totaling approximately 10,000 examples.

The primary dataset, krogoldAI/rag-query-analysis, contains 7,305 high-quality query-analysis pairs. This dataset contains queries carefully curated from three sources:

• rag-datasets/rag-mini-wikipedia
• razbit96/Ambiguity-Handling-in-User-Queries
• glaiveai/RAG-v1.

Approximately 20% of the training examples include queries with systematically introduced ambiguity at varying levels to ensure the model can handle realistic user inputs across the ambiguity spectrum.

The training data underwent rigorous quality assurance through a dual evaluation framework. Each example was validated for strict XML schema conformance and semantically evaluated using an LLM-as-a-judge protocol with six quality dimensions. Only examples achieving both perfect structural validity and high semantic quality scores were included in the final dataset, ensuring the model was trained exclusively on gold-standard examples.

Additionally, the model was trained on 3,000 randomly chosen examples of krogoldAI/rag-ambiguous-queries, a dataset of queries derived from glaiveai/RAG-v1 designed specifically for the GRPO phase. This dataset includes queries distributed across four synthetically induced ambiguity levels:

• no ambiguity (45%, original well-formed queries),
• low ambiguity (30%, minor vagueness from removing one specific detail),
• medium ambiguity (20%, moderate ambiguity with multiple missing references),
• high ambiguity (5%, heavy ambiguity with extensive pronoun usage).

This stratified distribution ensures the model learns to handle the full spectrum of query ambiguity while being trained primarily on realistic, naturally-occurring query patterns.

Training Procedure

QueryRefiner-0.5B-v0.1-GRPO underwent a three-phase training procedure combining reinforcement learning with two-stage supervised fine-tuning, all using full parameter updates (not parameter-efficient methods like LoRA) on Qwen2.5-0.5B-Instruct.

Phase 1: Reinforcement Learning with GRPO

The model was first trained using Group Relative Policy Optimization (GRPO) on 3,000 examples from krogoldAI/rag-ambiguous-queries for 1 epoch with a learning rate of 5e-6. This phase focused on learning correct XML structure and formatting before semantic refinement.

The GRPO reward function evaluated outputs through a weighted combination of five components: tag structure (30%), XML validity (25%), element ordering (25%), confidence formatting (18%), and confidence distribution (2%). The structure component verified the presence of all required XML elements, while validity ensured parseability. The ordering component checked that tags appeared in the correct sequence, and the confidence component validated that confidence values were properly formatted and summed to 1.0 for ambiguous cases.

The confidence distribution component used normalized entropy to encourage balanced probability distributions in ambiguous classifications, computed as:

$$ -\frac{1}{\log n}\sum_{j=1}^n p_j \log(p_j+\delta), $$

where \( (p_j)_{1\leqslant j\leqslant n} \) are the confidence scores and \( \delta = 10^{-8} \) is a small regularization constant. This penalizes overly skewed distributions (e.g., 0.99/0.01 splits) in favor of more balanced confidence allocations when multiple interpretations are plausible. Additional penalties were applied for structural issues (missing attributes) and for outputs that attempted to answer queries rather than analyze them.

The training code and the reward function code can be found here and here, respectively.

Phase 2: Breadth-Focused Supervised Fine-Tuning

Following structural learning through GRPO, the model underwent supervised fine-tuning on the complete krogoldAI/rag-query-analysis dataset (7,305 examples) for 2 epochs with a learning rate of 2e-5. This phase captured broad semantic patterns across diverse query types and domains.

The training code can be found here.

Phase 3: Depth-Focused Supervised Fine-Tuning

The final phase focused on precision refinement using only the highest-quality 3,660 examples from krogoldAI/rag-query-analysis, i.e., those examples where the judge LLM assigned perfect scores. This "polish" phase was trained for 3 epochs with a learning rate of 2e-5 to maximize learning from gold-standard examples.

This progressive training approach ensures the model achieves both reliable formatting and nuanced semantic understanding. All training was performed on an NVIDIA A100 SXM GPU.

The training code can be found here.

Model Capabilities

The model employs a systematic, ontology-inspired analysis framework that decomposes queries into structured XML representations. This analytical approach (which identifies domains, intents, concepts, relations, and ambiguities) aims to enhance the quality of the rephrased query output. For every query, it provides domain classification and intent detection with confidence scores that sum to 1.0, properly handling both unambiguous cases (single candidate with confidence 1.0) and ambiguous cases (multiple candidates with distributed confidence). The model can optionally extract explicit and implicit concepts, identify relations between entities using subject-predicate-object triples, and normalize ambiguous terms when disambiguation would improve retrieval.

The rephrasing capability focuses on retrieval optimization rather than query answering. The model transforms queries by using specific terminology likely to appear in relevant documents, expanding acronyms when contextually appropriate, adding disambiguating context, and making implicit references explicit through placeholder notation such as [PERSON] or [COMPANY]. Importantly, the model has learned to preserve already-optimal queries unchanged, recognizing when rephrasing would not improve retrieval effectiveness.

This structured analytical framework ensures that rephrasings are informed by comprehensive query understanding rather than surface-level transformations, leading to more semantically precise retrieval-optimized queries.

The default output schema is the following:

<query_analysis>
  <domain ambiguous="true|false">
    <candidate confidence="X.X">...</candidate>
  </domain>
  <intent ambiguous="true|false">
    <candidate confidence="X.X">...</candidate>
  </intent>
  <!-- Optional sections -->
  <concepts>
    <explicit>...</explicit>
    <implicit>...</implicit>
  </concepts>
  <relations>
    <relation subject="..." predicate="..." object="..."/>
  </relations>
  <normalized_terms>
    <term original="..." normalized="..."/>
  </normalized_terms>
  <!-- End optional sections -->
  <insufficient_context>true|false</insufficient_context>
  <rephrased>true|false</rephrased>
  <rephrased_query>...</rephrased_query>
</query_analysis>

Limitations and Considerations

As a 0.5B parameter model, QueryRefiner-0.5B-v0.1-GRPO prioritizes efficiency and deployability over the capabilities of larger language models. While it performs well on the types of queries represented in its training distribution, performance may degrade on highly specialized domains, multilingual queries, or query types significantly different from the training examples. The model focuses exclusively on English-language queries and has been optimized for the specific XML output format defined in its training.

The model's ambiguity detection and confidence scoring reflect patterns learned from the training data, which includes both natural and synthetically augmented ambiguous queries. While the training process incorporated diverse ambiguity levels, edge cases or novel forms of ambiguity may not be handled with the same reliability as more common patterns. Users should consider the model's confidence scores as informative signals rather than calibrated probabilities.

Usage Example

First, make sure you have the latest version of transformers:

pip install git+https://github.com/huggingface/transformers.git

Define the system prompt (since it was used as such during training, for optimal results we recommend not changing it).

<details> <summary><i>Unroll to see the system prompt</i></summary>

SYSTEM_PROMPT = """You are a query analysis and rephraser for a Retrieval-Augmented Generation (RAG) system.
Your sole task is to **analyze user queries** and output a structured XML document.
You must **not answer the query itself**, only analyze and rephrase it.

## RAG Query Optimization

Effective rephrasing should optimize for document retrieval by:
- Using **specific terminology** and domain vocabulary likely to appear in relevant documents
- **Expanding acronyms** when they add context (but not when the acronym itself is the subject)
- **Adding disambiguating context** without over-constraining the search
- **Making implicit references explicit** using placeholders for missing entities (e.g., [PERSON], [COMPANY])
- **Preserving user intent** while improving retrieval precision

Examples: "How do I reset my password?" → "password reset procedure authentication"
"What's their revenue?" → "What's [COMPANY]'s revenue?"

## Analysis Process

Follow this systematic approach to decompose each query:
1. **Identify the domain**: Determine the subject area or field the query relates to (e.g., banking, healthcare, technology, legal). Consider both explicit domain indicators and contextual clues.
2. **Determine the intent**: Classify what the user is trying to accomplish (e.g., definition lookup, troubleshooting, comparison, how-to guidance, factual question).
3. **Extract key concepts (optional)**: Identify explicit terms mentioned and relevant implicit concepts that would aid in query understanding.
4. **Identify relations (optional)**: Map out relationships between entities using subject-predicate-object triples when meaningful connections exist.
5. **Normalize terms (optional)**: Disambiguate or standardize ambiguous terms when clarification would improve retrieval (e.g., "Apple" → "Apple Inc." vs "apple fruit").
6. **Assess query quality**: Evaluate if the query has sufficient context for retrieval and whether rephrasing would improve it.
7. **Generate rephrased query**: Create a clearer, more specific version optimized for document retrieval, or keep the original if already optimal.

## Technical Rules

1. **Never answer the user's question.** Only analyze and rephrase.
2. Always produce valid XML strictly following the schema below.
3. `<domain>` and `<intent>` are **mandatory** and must contain one or more `<candidate confidence="X.X">...</candidate>` entries:
   - Confidence scores must always sum to 1.0
   - If unambiguous: **exactly one candidate** with `confidence="1.0"` and `ambiguous="false"`
   - If ambiguous: multiple candidates with `ambiguous="true"` and confidence distributed proportionally to plausibility:
     - Use uniform distribution only when candidates are genuinely equally likely
     - Otherwise, weight confidence toward the more probable interpretation
   - Examples:
     - "What is Mercury's rotation period?" → Astronomy 0.5, Chemistry 0.5 (equally plausible)
     - "Jaguar speed in the wild" → Zoology 0.8, Automotive 0.2 (context favors animal)
4. Confidence values must always have one decimal place (e.g., `0.5`, `1.0`).
5. Only `<concepts>`, `<relations>`, and `<normalized_terms>` are optional. **All other elements are mandatory.**
6. `<insufficient_context>` and `<rephrased>` must each appear **exactly once** and be either `true` or `false`.
7. `<rephrased_query>` must always appear, even if identical to the input.
8. **Output only valid XML.** Do not include any explanations, comments, or text outside the XML structure.
9. All elements must appear in the order specified in the schema:
   `<domain> → <intent> → <concepts> → <relations> → <normalized_terms> → <insufficient_context> → <rephrased> → <rephrased_query>`.

## Output Schema

<query_analysis>
  <domain ambiguous="true|false">
    <candidate confidence="X.X">...</candidate>
  </domain>
  <intent ambiguous="true|false">
    <candidate confidence="X.X">...</candidate>
  </intent>
  <!-- Optional sections -->
  <concepts>
    <explicit>...</explicit>
    <implicit>...</implicit>
  </concepts>
  <relations>
    <relation subject="..." predicate="..." object="..."/>
  </relations>
  <normalized_terms>
    <term original="..." normalized="..."/>
  </normalized_terms>
  <!-- End optional sections -->
  <insufficient_context>true|false</insufficient_context>
  <rephrased>true|false</rephrased>
  <rephrased_query>...</rephrased_query>
</query_analysis>
"""

</details>

Then, use the code below to get started with the model.

from transformers import AutoTokenizer, AutoModelForCausalLM

model_name = "krogoldAI/QueryRefiner-0.5B-v0.1-GRPO"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

user_query = "How do I reset my password?"

messages = [
    {"role": "system", "content": SYSTEM_PROMPT},
    {"role": "user", "content": user_query}
]

input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(input_text, return_tensors="pt")

outputs = model.generate(**inputs, temperature=0.7, max_new_tokens=512)
analysis = tokenizer.decode(outputs[0], skip_special_tokens=True)

print(analysis)

Evaluation and Performance

Structural Validation

To assess the model's ability to produce correctly formatted outputs, we evaluated structural conformance across 1,000 examples from the test split of krogoldAI/rag-ambiguous-queries, comparing both QueryRefiner-0.5B-v0.1 models against their base model Qwen2.5-0.5B-Instruct. The evaluation measured adherence to the required XML schema, including tag presence, well-formedness, element ordering, and confidence score formatting.

Table 1 - Structural validity metrics (% of outputs meeting each requirement)

Here, tag structure verifies that all required XML tags are present, XML validity ensures the output is well-formed and parseable, order confirms that required tags appear in the correct sequence, and confidence validates that confidence values are properly formatted and sum to 1.0.

<!-- The results demonstrate that fine-tuning substantially improved structural reliability, with QueryRefiner achieving near-perfect conformance (99.6%) compared to the base model's inability to produce any fully compliant outputs. -->

Semantic Validation

Beyond structural correctness, we evaluated the semantic quality of the model's outputs using an LLM-as-a-judge protocol with GPT5 on 1,000 examples from the test split of krogoldAI/rag-ambiguous-queries. Each output was assessed across six dimensions aligned with the model's core objectives.

Table 2 - Semantic scores

All values are reported as mean ± standard deviation (%), computed over test examples. The base model produced too few valid XML samples for meaningful semantic evaluation.

<details> <summary><i>Unroll to see the system prompt used for the "judge" LLM</i></summary>

JUDGE_PROMPT = """You are evaluating query analyses for a RAG system.

### System Requirements
The analyzer was instructed to optimize queries for document retrieval by:
- Using **specific terminology** and domain vocabulary likely to appear in relevant documents
- **Expanding acronyms** when they add context (but not when the acronym itself is the subject)
- **Adding disambiguating context** without over-constraining the search
- **Making implicit references explicit** using placeholders for missing entities (e.g., [PERSON], [COMPANY])
- **Preserving user intent** while improving retrieval precision
- **Keeping the original query unchanged** if it's already well-optimized for retrieval

### Input
Original: "{original}"
Domain: {domain}
Intent: {intent}
Rephrased: "{rephrased}"

Note: The [ambiguous] tag indicates the analyzer determined the query has multiple plausible interpretations for that dimension, with confidence distributed across candidates.

### Evaluation Criteria (1-5 scale)

1. Domain Accuracy (1=wrong, 3=acceptable, 5=perfect)
   - Are the domain candidates correct?
   - Are confidence scores reasonable?

2. Intent Accuracy (1=wrong, 3=acceptable, 5=perfect)
   - Are the intent candidates correct?
   - Are confidence scores reasonable?

3. Ambiguity Assessment (1=wrong, 3=acceptable, 5=perfect)
   - Is the ambiguity determination appropriate for this query?
   - If ambiguous: Is the confidence distribution justified?
   - If clearly unambiguous but marked ambiguous (or vice versa), score ≤2.

4. Rephrasing Quality
   1 = Poor (significantly degraded the query, or completely failed to address clear issues)
   2 = Suboptimal (minor degradation, or missed an obvious improvement opportunity)
   3 = Neutral (minor changes with mixed effects)
   4 = Good improvement, but could be better
   5 = Optimal outcome (either improved a suboptimal query, or correctly preserved an already-optimal one)
   (Note: Do not penalize rephrasing for being minimal if the original was already optimal.)

5. Intent Preservation (1=lost, 3=mostly preserved, 5=fully preserved)
   - Focus on meaning fidelity, not retrieval optimization.

6. Follows Guidelines (1=violates, 3=mostly follows, 5=perfectly follows)
   - Check adherence to the RAG optimization principles above.

### Output Format
{{
  "domain_accuracy": <1-5>,
  "intent_accuracy": <1-5>,
  "ambiguity_assessment": <1-5>,
  "rephrasing_quality": <1-5>,
  "intent_preservation": <1-5>,
  "follows_guidelines": <1-5>,
  "critical_issue": "<brief description or null>",
  "usable": <true/false> // true if suitable for RAG use, even if not perfect
}}

Output only valid JSON. Do not include any explanations, comments, or text outside the JSON structure.
"""

</details> <!-- To further examine the model's robustness across varying query complexity, we analyzed performance stratified by ambiguity level. The following table shows results broken down by the four ambiguity categories in the test set (NONE: 45%, LOW: 30%, MEDIUM: 20%, HIGH: 5%), which reflect the distribution used during the GRPO training phase. --> <!-- #### Table 3 - Per-ambiguity breakdown --> <!-- | | NONE | LOW | MEDIUM | HIGH | |:---------------------|:---------------|:---------------|:---------------|:---------------| | domain_accuracy | 98.82 ± 9.48% | 98.91 ± 8.53% | **99.26 ± 6.54%** | 99.04 ± 4.85% | | intent_accuracy | 98.77 ± 8.94% | **99.24 ± 6.93%** | 98.77 ± 6.46% | 95.67 ± 10.76% | | ambiguity_assessment | 99.27 ± 6.56% | 99.33 ± 6.46% | 99.63 ± 3.02% | **100.00 ± 0.00%** | | rephrasing_quality | **90.81 ± 17.49%** | 89.98 ± 18.80% | 90.52 ± 15.11% | 87.50 ± 20.71% | | intent_preservation | 96.13 ± 14.72% | 95.88 ± 15.71% | **96.80 ± 12.02%** | 92.79 ± 18.08% | | follows_guidelines | 97.14 ± 13.26% | 97.05 ± 12.91% | **98.89 ± 7.17%** | 95.67 ± 13.76% | | semantic_score | 96.82 ± 9.58% | 96.73 ± 9.80% | **97.31 ± 6.82%** | 95.11 ± 10.06% | --> <!-- The evaluation results reveal a clear task difficulty hierarchy that aligns with the inherent complexity of each component. Structural and classification metrics (domain accuracy, intent accuracy, ambiguity assessment) achieve 98-99% performance, while the generative rephrasing task scores lower at ~90%. This gap reflects the fundamental difference in task complexity rather than a training deficiency. Domain and intent classification are essentially pattern recognition tasks where the model must map queries to learned categories, a task well-suited to the model's 0.5B parameter capacity. Similarly, ambiguity assessment and guideline adherence involve rule-following and structural analysis, which the three-phase training procedure was explicitly designed to optimize. Rephrasing quality, by contrast, requires the model to make nuanced judgments about when to intervene (versus preserving already-optimal queries), how to balance specificity against over-constraint, and whether to expand acronyms or add disambiguating context. These decisions demand deeper semantic understanding and generation capabilities that push the limits of a 0.5B model. The 90.33% score with 17.61% standard deviation represents strong performance on this challenging task, particularly given the diversity of the test set. For production deployments, users should anticipate that the model will perform most reliably on structural conformance and classification tasks, while rephrasing decisions may occasionally require human review, particularly for edge cases or highly specialized domains underrepresented in the training data.-->

Performance Considerations

Performance characteristics will vary based on query type, domain, and ambiguity level. The model is expected to perform strongest on queries similar to those in the training distribution and may require additional fine-tuning or prompt engineering for specialized applications or domains underrepresented in the training data.

Acknowledgments

This model builds upon Qwen2.5-0.5B-Instruct by the Qwen team at Alibaba Cloud. The training data incorporates queries from rag-datasets/rag-mini-wikipedia, razbit96/Ambiguity-Handling-in-User-Queries, and glaiveai/RAG-v1.