IbrahimKhan7208/investment-research-router

Investment Research Bot — Router Model

A small, fine-tuned Qwen2.5-1.5B-Instruct model that replaces two LLM calls in a LangGraph-based multi-tool investment research agent (classifierNode + extractFiltersFromQuestions) with a single structured-output call — trained through a full CPT → SFT → GRPO pipeline.

Why this exists

The original agent used two separate LLM calls per research question, both served by a 120B-class model:

1. 1.`classifierNode` — decomposes the question into sub-questions and assigns each a tool (RAG, WEB, or STOCK).
2. 2.`extractFiltersFromQuestions` — for each RAG sub-question, a second LLM call extracts companies, years, and a cleaned search query. WEB/STOCK sub-questions got no structured extraction at all — the raw question string was passed straight to Tavily / Yahoo Finance.

This also hardcoded the RAG-eligible company list to exactly three companies (NVIDIA, AMD, Microsoft), baked directly into both prompts.

The router model collapses this into one model call that outputs tool routing and structured extraction (companies, years, search query) for all three tools uniformly — including WEB/STOCK, which the original system never structured at all. It also drops the hardcoded company enum: extraction generalizes to any named company, not just the original three.

Pipeline

Base model (Qwen2.5-1.5B-Instruct)
       |
Continued Pretraining (CPT)   -- domain vocabulary/register adaptation
       |
Supervised Fine-Tuning (SFT)  -- teaches the routing/extraction task
       |
Reinforcement Learning (GRPO) -- closes schema/edge-case gaps SFT couldn't
       |
grpo_router_final_merged  <- final deployed router

1. Base model

Qwen2.5-1.5B-Instruct — chosen for its small footprint (deployable on modest hardware, fast single-call inference) relative to the 120B-class model used for the original two-call baseline. The Instruct checkpoint was kept throughout the entire pipeline (never switched to the raw base model) — a deliberate choice: the SFT dataset was sized assuming the model already knows how to follow instructions and only needs to specialize, not learn instruction- following from scratch. CPT was correspondingly kept light-touch, specifically because it perturbs an already-aligned model.

2. Continued Pretraining (CPT)

Goal: light domain-vocabulary adaptation to SEC-filing register (financial/legal terminology, disclosure structure) before task-specific SFT — not deep domain rewiring, and not meant to teach the model to answer financial questions (that's SFT's job).

Corpus: TheFinAI/SEC_2025 only — real FY2025 SEC filing text (8-K item disclosures + 10-Q MD&A excerpts). A second candidate corpus (a scraped tech-news dataset) was evaluated and deliberately dropped in favor of a single-source, controlled experiment — avoids confounding CPT's effect with an unrelated data-quality/licensing variable. A mixed-source news+filings CPT experiment remains a documented future option, not something ruled out permanently.

Final corpus: 14 deduplicated documents, ~6,000 tokens (Qwen2.5 tokenizer), capped at 600 tokens/doc. Corpus construction went through two iterations — the first pull (10 docs) surfaced zero mentions of the target priority companies (NVIDIA, AMD, Microsoft, Apple, etc.) despite the source dataset containing them; the builder was hardened with explicit stream shuffling, per-row progress logging, and a scan cap before the second pull, which surfaced 3 genuine filings by priority companies (AMD's own executive-appointment 8-K, Salesforce's and Amazon's own earnings 8-Ks) among the 14. The corpus skews toward Item 5.02 personnel-filing boilerplate (8 of 14 docs) over Item 2.02/MD&A financial-results language (6 of 14) — a known, accepted trade-off given the light-touch CPT scope, not a blocking issue.

Training config (Unsloth + LoRA, QLoRA/4-bit):

• —r=16, lora_alpha=32 (2:1 ratio — CPT convention, more headroom than SFT's later 1:1 to shift the embedding-space distribution)
• —target_modules: attention + FFN plus `embed_tokens` and `lm_head` — the key CPT-specific inclusion, since vocabulary/output-distribution adaptation is the whole point of this stage; SFT's LoRA later deliberately excludes these two modules
• —learning_rate=2e-4, embedding_learning_rate=2e-5 (10x smaller — embeddings are the model's foundation, updated gently)
• —3 epochs (more than a typical CPT pass would need for a larger corpus — needed here specifically because the corpus is tiny; repeating a small corpus more times is standard practice and doesn't meaningfully raise catastrophic-forgetting risk given the corpus size stays small)
• —packing=True, optim=adamw_8bit, lr_scheduler_type=cosine
• —No 80/20 domain/general-text data mixing — judged unnecessary at this scale (tiny corpus, few epochs, LoRA-only, frozen base weights); a direct instruction-following sanity check was used instead of mixing as the catastrophic-forgetting safeguard (see below)

Evaluation — two checks, not just perplexity, since CPT here perturbs an aligned Instruct model and instruction-following damage was the specific risk being guarded against:

The perplexity drop is intentionally modest (Qwen2.5 already has substantial financial text in its pretraining mix, so there wasn't much "surprise" left to remove — a large drop would have suggested overfitting the tiny corpus, not healthy adaptation). The sanity- generation check passing cleanly confirmed the higher-risk part of this stage — embed_tokens/lm_head perturbation on an aligned model — didn't degrade instruction-following. Held-out set size (4 docs) means the perplexity number is directional, not precise.

Output: cpt_router_adapter (LoRA adapter, not merged) → later merged into a standalone checkpoint, cpt_merged_model, which SFT was built on top of.

3. Supervised Fine-Tuning (SFT)

3.1 Dataset construction

Synthetic dataset generated by a custom pipeline (generate_dataset.py

• —builders.py + assemble.py), targeting a fixed JSON output schema:

{
  "subQuestions": [
    {
      "question": "string",
      "tool": "RAG" | "WEB" | "STOCK",
      "companies": ["string", ...],
      "years": [int, ...],
      "searchQuery": "string (RAG only)"
    }
  ],
  "requiredTools": ["RAG", "WEB", "STOCK"]
}

Key design principles (see DATASET_CARD.md for full rationale):

• —No hardcoded company/year enum — companies are extracted as a general skill, not memorized from a fixed list.
• —`years: []` for unresolved relative time ("recent", "latest", "last year") — resolving relative time to concrete years is left to the app layer, not baked into the model.
• —Minimal splitting — multiple companies/metrics for the same tool stay in ONE sub-question; only a genuine tool change forces a split.
• —Canonical-name extraction — tickers (NVDA) and known aliases (Facebook) resolve to the canonical company name (NVIDIA, Meta) in the output, never the raw form as written.

3.2 Dataset scaling

v3 additions over v2: expanded entity vocab (finance, retail, energy, healthcare, telecom, airlines, industrials, international tech — 64 companies total), a ticker-to-canonical-name map (40 tickers), register diversity (informal abbreviations, Slack-terse phrasing, formal analyst-memo phrasing), 7 boundary-trap builders (RAG/WEB/STOCK boundary questions designed to be hard), multi-hop dependent-phrasing builders, and new structural shapes (reversed tool orderings, mixed company-count-bucket compounds) — raising distinct structural signature coverage from 10 to 25, specifically to avoid scaling by permuting entities within the same handful of shapes.

3.3 Post-hoc dataset patch (v3.1) — attempted, then deliberately not used

After the initial SFT run, out-of-domain evaluation surfaced 4 real gaps: alias-name canonicalization not generalizing (tickers worked, plain-English aliases like "Facebook" didn't), a schema leak (searchQuery appearing on STOCK sub-questions with explicit historical years), relative-time phrasing beyond "recent/latest/current" not resolving to years: [], and a generator bug in one compound- question builder where the target year/period wasn't stated in the visible question text.

A small 374-example patch dataset was built and used for a short continued-SFT pass. Two separate patch attempts (4 epochs, then 2 epochs on a clean-restarted model) both introduced new regressions on the original held-out val set — schema-field hallucinations, requiredTools-ordering errors — while only partially closing the original gaps. This pointed to instability from training on too narrow/isolated a correction set rather than a fixable epoch-count issue.

Decision: the patch was discarded. The final SFT checkpoint used downstream is the original v3, unpatched, 8,513-example SFT run. The 4 known gaps were deliberately deferred to the GRPO stage instead, where they could be addressed via reward-function checks on live rollouts rather than more small-scale SFT correction passes.

3.4 SFT training run

• —Framework: Unsloth + TRL SFTTrainer, QLoRA-style LoRA adapters, starting from cpt_merged_model (the CPT-adapted checkpoint, not raw base) — target_modules restricted to attention + FFN only this time (no embed_tokens/lm_head — that was CPT's job, already merged in)
• —r=16, lora_alpha=16 (1:1 — standard SFT convention, behavior shaping through existing representations rather than pushing the embedding space)
• —Trainable parameters: 18,464,768 of 1,562,179,072 (1.18%)
• —3 epochs, effective batch size 16 (batch 4 x grad accumulation 4, later switched to packing=True mid-run for throughput), response-only loss masking (train_on_responses_only — loss computed only on the assistant's JSON output, not the repeated system prompt/question)
• —Final training loss: 0.000626, final validation loss: 0.001455 — monotonically improving through the full run, no overfitting reversal

3.5 SFT evaluation

On a 120-question held-out sample from the v3 val set:

The single mismatch was later confirmed to be a dataset artifact (a compound-question builder not stating its target year/period in the visible question), not a model error — see §3.3. Out-of-domain checks (generic non-router prompts) confirmed no mode collapse into always-emit-JSON behavior.

Output: sft_router_merged — the checkpoint GRPO starts from.

4. Reinforcement Learning (GRPO)

Starting checkpoint: sft_router_merged (the original, unpatched v3 SFT model — not the abandoned v3.1 patch, see §3.3).

Why GRPO, and why these specific gaps: the 4 gaps deferred from SFT (§3.3) are all programmatically verifiable from the question text or against a gold label — exactly GRPO's ideal use case. A reward function built on rule-based checks can directly shape these behaviors via live rollouts, rather than depending on a small, narrow correction dataset (which SFT patching had already shown to be unstable).

4.1 Reward function

router_reward.py is fully rule-based — no separate reward model. It combines a hard schema-validity gate (malformed JSON or a schema violation like searchQuery on a non-RAG sub-question caps the reward at a fixed penalty, regardless of anything else) with weighted component checks:

Most components are self-checkable from the question alone (no gold label needed) — the GRPO prompt pool for those checks can be larger than the labeled SFT set. The reward function's ticker/alias list was deliberately kept simple: single-letter tickers ("V" for Visa, "T" for AT&T) and tickers that double as ordinary English words (NOW, SHOP, COST, SNOW) were excluded rather than special-cased, since real usage essentially never writes a bare single letter or common word as a ticker reference. Word-boundary matching (not naive substring matching) and an explicit-year regex covering both "in 2022" and "FY2024" phrasing were fixed after an initial version produced false positives on real model output — verified via a 15-case self-test and manual inspection against live generations before being trusted for training.

4.2 Prompt pool

2,198 questions total, three sources:

4.3 Training attempts

Three training attempts were run; the middle one produced the final checkpoint.

Attempt 1 — failed, no meaningful learning. Initial config (num_generations=6, default temperature, learning_rate=5e-6, 400 planned steps) was first found to be too slow to be practical on a free-tier T4 (~0.14 it/s); an attempt to enable Unsloth's vLLM-backed fast-generation path failed due to a CUDA version mismatch (vLLM built for CUDA 13, Colab's runtime on CUDA 12.8) and was abandoned in favor of trimming the standard-generation config instead (num_generations=4, max_completion_length=160, max_steps=200). That trimmed run completed, with misleadingly high aggregate reward (mostly 2.0-3.0 out of 3.0) — but reward_std was 0.000000 on over half of logged steps, and kl divergence never meaningfully exceeded ~0.002 across the entire run. Diagnosis: the SFT-converged policy was confident enough (SFT final loss ~0.0006) that default sampling produced near-identical completions across the group on most prompts — no differentiated signal for GRPO to learn from. Post-training evaluation confirmed all 4 target gaps were completely unchanged from the pre-GRPO baseline, including on the exact reward components they were meant to fix. Key lesson: a high aggregate reward number can fully mask zero progress on the specific failures a reward function was designed to catch — always check `reward_std`/`kl`, and evaluate against the specific target cases, not just the mean reward.

Attempt 2 — the final checkpoint. Restarted from sft_router_merged fresh (attempt 1's near-zero policy movement meant nothing was lost by not continuing from it). Two changes targeted the diagnosed cause directly: temperature=1.3, top_p=0.95 (forces real sampling diversity) and learning_rate=1e-5 (up from 5e-6 — headroom to move when signal exists). Prompt pool reweighted toward the harder gap/noise subsets (300 gold + all gap-targeted x2 + all noise x2 = 1,094 questions) rather than the full 2,198, concentrating training budget on what needed fixing. kl divergence reached as high as 0.0248 at individual steps (vs. attempt 1's ~0.002 ceiling) and reward_std showed real, sometimes large, spread (e.g. 0.948 at step 30) — genuine differentiated learning signal this time. Trained to step 175 (of a 200-step plan; the runtime session ended before completion, and the checkpoint-175 save was used as-is rather than continuing).

Evaluation at checkpoint-175:

Regression check against the original 8,513-question SFT benchmark: parse 100/100, schema 100/100, tool-sequence 100/100, exact match 99/100 — statistically identical to the pre-GRPO SFT baseline. Net result: real, generalizing improvement on one gap, a safety-relevant partial fix on a second, zero cost to the base task.

Attempt 3 — a targeted top-up, discarded. Continued training directly from the attempt-2 checkpoint (not a fresh LoRA) on a pool concentrated almost entirely on the one remaining clean-failure gap (alias/ticker canonicalization weighted 3x, explicit-year 2x, only a small 120-question gold anchor for regression protection — 637 questions total, 80 steps). Result: no improvement on alias canonicalization (still 0.0 on every tested case) and a net regression — a previously-correct compound RAG question now failed schema validation (searchQuery missing), and the regression-suite exact-match score dropped from 99/100 to 91/100. Discarded; attempt 2's checkpoint-175 remains the final model. Alias/name-based canonicalization has now failed to improve across four separate correction attempts (two SFT patches, one general GRPO run, one targeted GRPO top-up) — treated as a documented, deliberate open limitation rather than a target for further correction attempts within this project's scope (see §7).

4.4 Final merge

sft_router_merged (base) + attempt 2's checkpoint-175 adapter -> merged -> `grpo_router_final_merged`, the final deployed router. Post-merge sanity check (this project's own canonical example question: "Compare NVIDIA vs AMD's data center revenue, recent news sentiment, and stock performance") produced fully correct, schema-valid, correctly-decomposed output on the first try.

5. Benchmark: Router vs. Original Baseline

Baseline = original classifierNode + extractFiltersFromQuestions prompts, reproduced verbatim, served via Groq (openai/gpt-oss-120b), JSON mode forced. Router = grpo_router_final_merged, run unbatched on a free-tier Colab T4.

Evaluated on two held-out subsets, split deliberately to keep the comparison honest:

• —In-scope: questions whose companies are within the baseline's original hardcoded enum (NVIDIA / AMD / Microsoft) — a genuine apples-to-apples routing comparison.
• —Out-of-scope: questions naming companies outside that enum — the baseline cannot serve these by design (the enum is hardcoded into its prompt), so this subset demonstrates a capability gain, not a routing-accuracy contest.

In-scope (n=31, fair fight: NVIDIA/AMD/MSFT only)

Out-of-scope (n=30, baseline's enum can't handle these)

*\Latency caveat: baseline runs on Groq's hosted inference infrastructure; the router ran unbatched on a free-tier T4 with no serving optimization. This gap reflects infrastructure, not model quality** — it is not presented as a head-to-head speed claim. With proper serving (e.g. vLLM, batching, or hosted inference), the router's 1.5B parameter count and single-call design would be expected to be substantially faster than the 120B two-call baseline in practice.

Reading these numbers honestly

• —In-scope accuracy gap (61.3% -> 96.8%): the val set specifically includes RAG/WEB/STOCK "boundary trap" questions (e.g. a filing metric using STOCK-suggestive wording like "shares outstanding") designed to be hard. The original baseline prompt has no rules for these boundary cases; the router's training data was built specifically to cover them.
• —In-scope RAG field accuracy gap (41.2% -> 94.1%): partly a genuine design-philosophy difference, not purely "worse extraction." The baseline's filter-extraction prompt explicitly resolves relative time ("recent", "latest") to concrete years ([2024, 2025]); the router's schema requires years: [] for unresolved relative time. Every relative-time RAG question is a near-automatic baseline "miss" under the router's schema by construction — still a fair comparison since the router's schema is what the downstream system actually consumes, but worth stating plainly rather than implying the baseline extracts sloppily across the board.
• —Out-of-scope numbers: not a baseline failure so much as a scope mismatch — the original prompt hardcodes exactly 3 companies. The router's 87.5-100% here demonstrates genuine generalization (extraction works on any company, no code changes, no retraining), which is the headline capability claim of this project.
• —Sample sizes (n=~30 per subset) are enough to be directionally solid but small enough that any single percentage point (e.g. the out-of-scope subset slightly outperforming in-scope) shouldn't be over-read — treat these as illustrative, not a rigorous statistical claim, unless re-run at a larger n.

6. Repo structure

generate_dataset.py       # system prompt, vocab lists, ticker/alias maps
builders.py                # dataset builder functions + hand-crafted edge cases
assemble.py                 # weighted sampling, dedup, train/val split, stats
router_sft_train.jsonl      # 7,662 SFT examples
router_sft_val.jsonl        # 851 SFT examples
DATASET_CARD.md              # full schema + design rationale + composition

build_cpt_corpus_colab.py   # CPT corpus builder (TheFinAI/SEC_2025, stream-shuffled, deduplicated)
cpt_notebook.py               # CPT training notebook (Qwen2.5-1.5B-Instruct, QLoRA)

sft_notebook.py                # SFT training notebook (chat-template formatting, response-only loss masking)

router_reward.py                 # GRPO reward function (schema + rule-based checks)
grpo_prompt_pool.jsonl            # 2,198-question GRPO prompt pool (gold + gap-targeted + noise)

grpo_router_final_merged/          # final deployed model

7. Known limitations

• —Alias/name-based company canonicalization does not work. "Facebook" does not resolve to "Meta," "Alphabet" does not resolve to "Google," and similar plain-English aliases are extracted verbatim rather than canonicalized. Ticker-based resolution (NVDA -> NVIDIA, MSFT -> Microsoft) works correctly and generalizes well. This gap was targeted across 4 separate correction attempts (2 SFT patch passes, 1 general GRPO run, 1 targeted GRPO top-up) — none succeeded, and the final targeted attempt actively regressed the base task instead. Documented here as a deliberate, accepted limitation rather than continuing to spend training budget on diminishing returns; a dedicated, isolated SFT pass with much heavier alias- example density is the most likely path to actually closing this, left as future work.
• —Explicit-year extraction on STOCK/WEB sub-questions is incomplete. The dangerous half of this gap (an illegal searchQuery field appearing on a STOCK/WEB sub-question, which would break downstream JSON.parse()-based consumption) is fixed. The correctness half (populating years with the stated year, e.g. "Tesla's stock price in 2022" -> years: [2022]) is not — the model currently returns years: [] in these cases instead.
• —Latency benchmark reflects unoptimized single-request inference on a free-tier GPU, not a production serving setup — see §5 caveat.
• —Benchmark sample sizes (~30/subset) are illustrative; a larger-n re-run is recommended before quoting these percentages in a more formal context.
• —CPT's held-out perplexity evaluation used only 4 documents — directionally useful, not a precise measurement.

8. Acknowledgments

Built with Unsloth, TRL, and Groq. Base model: Qwen2.5-1.5B-Instruct (Alibaba Cloud / Qwen team). CPT domain corpus: TheFinAI/SEC_2025.