Update README.md

README.md

@@ -19,7 +19,7 @@ base_model:
 
 QueryRefiner-0.5B-v0.1-GRPO is a specialized query analysis and rephrasing model fine-tuned from [Qwen2.5-0.5B-Instruct](https://huggingface.co/Qwen/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 addresses 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. -->
+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.
 
@@ -131,7 +131,9 @@ Typical deployment scenarios include conversational search systems, question-ans
 
 ## Training Data
 
-The model was trained on the [krogoldAI/rag-query-analysis](https://huggingface.co/datasets/krogoldAI/rag-query-analysis) dataset, which contains 7,305 high-quality query-analysis pairs. This dataset was carefully curated from three sources:
+The model was trained on two complementary datasets totaling approximately 10,000 examples.
+
+The primary dataset, [krogoldAI/rag-query-analysis](https://huggingface.co/datasets/krogoldAI/rag-query-analysis), contains 7,305 high-quality query-analysis pairs carefully curated from three sources:
 - [rag-datasets/rag-mini-wikipedia](https://huggingface.co/datasets/rag-datasets/rag-mini-wikipedia)
 - [razbit96/Ambiguity-Handling-in-User-Queries](https://huggingface.co/datasets/razbit96/Ambiguity-Handling-in-User-Queries)
 - [glaiveai/RAG-v1](https://huggingface.co/datasets/glaiveai/RAG-v1).
@@ -140,49 +142,37 @@ Approximately 20% of the training examples include queries with systematically i
 
 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](https://huggingface.co/datasets/krogoldAI/rag-ambiguous-queries), a dataset of queries derived designed specifically for the GRPO phase. This dataset includes queries distributed across four 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
 
-The model underwent full fine-tuning (not parameter-efficient methods like LoRA) of all parameters in [Qwen2.5-0.5B-Instruct](https://huggingface.co/Qwen/Qwen2.5-0.5B-Instruct). Training was conducted over three epochs with a per-device batch size of 4 and gradient accumulation over 4 steps, yielding an effective batch size of 16. The learning rate was set to 2e-5 with a weight decay of 0.01 to prevent overfitting. A warmup period of 100 steps allowed the model to gradually adapt to the task-specific distribution before reaching the full learning rate. Training was performed on an NVIDIA A100 SXM GPU.
+QueryRefiner-0.5B-v0.1-GRPO underwent a sophisticated 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](https://huggingface.co/Qwen/Qwen2.5-0.5B-Instruct).
 
-The training code can be found [here](Code/Fine-tuning.py).
+### 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](https://huggingface.co/datasets/krogoldAI/rag-ambiguous-queries) for 1 epoch with a learning rate of 5e-6 (intentionally low for GRPO stability). This phase focused on learning correct XML structure and formatting before semantic refinement.
 
-GRPO on 3,000 data from [krogoldAI/rag-ambiguous-queries](https://huggingface.co/datasets/krogoldAI/rag-ambiguous-queries/) with:
+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.
 
-```py
-NUM_EPOCHS = 1
-PER_DEVICE_BATCH = 8
-GRADIENT_ACCUMULATION = 1
-LEARNING_RATE = 5e-6 # Low for GRPO stability
-NUM_GENERATIONS = 2 # Number of generations per prompt for GRPO
-MAX_NEW_TOKENS = 512  # Sufficient for XML output
-WARMUP_STEPS = 50
-```
+The confidence distribution component used normalized entropy to encourage balanced probability distributions in ambiguous classifications, computed as:
 
-SFT on 7,305 data from [krogoldAI/rag-query-analysis](https://huggingface.co/datasets/krogoldAI/rag-query-analysis/) with:
+$-\frac{1}{\log n}\sum_{j=1}^n p_j \log p_j$
 
-```py
-NUM_EPOCHS = 2
-PER_DEVICE_BATCH = 4
-GRADIENT_ACCUMULATION = 4
-LEARNING_RATE = 2e-5
-WEIGHT_DECAY = 0.01
-WARMUP_STEPS = 100
-```
+where \\( (p_j)_{1\leq j\leq n} \\) are the confidence scores, \\( n \\) is the number of confidence scores, and \\( \log n \\) represents the maximum possible entropy. This penalizes overly skewed distributions (e.g., 0.9/0.1 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.
 
-(--> Breadth)
+### 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](https://huggingface.co/datasets/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.
 
-Second SFT on 3,660 examples with:
+### 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](https://huggingface.co/datasets/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.
 
-```py
-NUM_EPOCHS = 3
-PER_DEVICE_BATCH = 4
-GRADIENT_ACCUMULATION = 4
-LEARNING_RATE = 2e-5
-WEIGHT_DECAY = 0.01
-WARMUP_STEPS = 50
-```
+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.
 
-(--> Depth)
 
 ## Model Capabilities
 
@@ -348,7 +338,7 @@ print(analysis)
 
 ### 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](https://huggingface.co/datasets/krogoldAI/rag-ambiguous-queries), comparing QueryRefiner-0.5B-v0.1 against its base model [Qwen2.5-0.5B-Instruct](https://huggingface.co/Qwen/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.
+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](https://huggingface.co/datasets/krogoldAI/rag-ambiguous-queries), comparing both QueryRefiner-0.5B-v0.1 models against their base model [Qwen2.5-0.5B-Instruct](https://huggingface.co/Qwen/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)
 
@@ -399,8 +389,6 @@ Beyond structural correctness, we evaluated the semantic quality of the model's
 |Overall semantic score|96.10 ± 9.74%|96.45 ± 9.52%|**96.81 ± 9.18%**|
  -->
 
-(Corrected for the SFT + GRPO-SFT models)
-
 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>