Sepehr Karimi, Ph.D.

Sepehr Karimi, Ph.D.

Director of AI and IT Solutions

HydroUA-GPT: Domain-Specific Hydrology LLM Fine-Tuning Pipeline

HydroUA-GPT hydrology LLM fine-tuning pipeline thumbnail

Overview

I led the development of HydroUA-GPT, a domain-specific hydrology large language model pipeline designed to turn hydrology textbooks, lecture material, and recent research literature into a reproducible AI training workflow.

The project focuses on the full AI engineering lifecycle for a scientific LLM: corpus extraction, data cleaning, synthetic instruction-data generation, parameter-efficient fine-tuning, inference, optional retrieval-augmented generation, and GPU/HPC deployment. The goal was not simply to build a chatbot, but to create a maintainable training and deployment pipeline for adapting Llama-family models to hydrology and water-resources knowledge.

My Role

  • Led the design and development of the HydroUA-GPT codebase and training workflow.
  • Refactored a messy collection of one-off scripts into a modular Python package with clean CLI entry points.
  • Designed the end-to-end LLM pipeline from raw PDF sources to fine-tuned hydrology assistant.
  • Implemented configurable corpus extraction, QA generation, fine-tuning, inference, RAG indexing, and Gradio deployment modules.
  • Built GPU/HPC-ready scripts for large-model QA generation and parameter-efficient fine-tuning.
  • Structured the project so sensitive assets such as PDFs, generated datasets, checkpoints, adapters, and vector stores remain outside GitHub.

AI Engineering Focus

HydroUA-GPT is organized around a reproducible model-adaptation workflow:

  1. Hydrology corpus construction from source PDFs using PyMuPDF-based extraction, paragraph segmentation, whitespace normalization, source metadata, page tracking, and duplicate filtering.
  2. Synthetic instruction-data generation using an instruction-tuned LLM through vLLM, with tensor parallelism, batched generation, deterministic sampling, restart/resume support, and QA parsing.
  3. Instruction dataset preparation that converts generated QA records into JSONL rows with instruction, input, output, source, and page metadata.
  4. LoRA/QLoRA supervised fine-tuning over Llama-family models with trainable adapter layers instead of full-model weight updates.
  5. Hydrology assistant inference using Transformers/PEFT loading, optional adapter attachment, 4-bit loading support, and a hydrology-specific system prompt.
  6. Optional RAG extension using SentenceTransformers embeddings and FAISS indexing over recent hydrology papers to provide retrieval context at inference time.
  7. Application deployment through a lightweight Gradio interface for interactive hydrology Q&A.

Fine-Tuning Method

The implemented fine-tuning workflow uses parameter-efficient supervised instruction tuning. Instead of updating every model weight, HydroUA-GPT trains LoRA adapters on targeted transformer modules:

  • attention projections: q_proj, k_proj, v_proj, o_proj
  • feed-forward projections: gate_proj, up_proj, down_proj

This is a stronger engineering choice than naive full fine-tuning for this project because it reduces GPU memory requirements, makes experiments cheaper and faster, keeps the base model reusable, and allows multiple hydrology-domain adapters to be versioned independently.

The training implementation includes several practical details important for reliable LLM adaptation:

  • Prompt/response loss masking: prompt tokens are masked with -100, so training loss is applied only to the assistant response rather than teaching the model to reproduce the user prompt.
  • Chat-template formatting: hydrology questions are converted into Llama-style system/user/assistant conversations.
  • bfloat16 training: the default training path uses BF16 for modern NVIDIA GPUs.
  • Gradient checkpointing: memory usage is reduced during backpropagation.
  • Gradient accumulation: small per-device batches are accumulated into a larger effective batch size.
  • Cosine learning-rate schedule with warmup: training is stabilized while adapting to the hydrology domain.
  • Train/evaluation split: a held-out validation subset can be used for step-based evaluation.
  • QLoRA option: the same pipeline can load the base model in 4-bit NF4 quantization for lower-memory training.

System Architecture

HydroUA-GPT is structured as a maintainable Python package with config-driven modules:

  • configs/books.yaml — source PDF registry and extraction settings.
  • configs/qa_generation.yaml — vLLM model path, tensor parallelism, sampling, and QA generation settings.
  • configs/finetune_lora.yaml — base model, dataset path, LoRA parameters, batch sizes, scheduler, precision, and checkpointing.
  • configs/chat.yaml — inference model path, adapter path, system prompt, decoding settings, and quantization options.
  • configs/rag.yaml — paper metadata path, FAISS index path, embedding model, and retrieval settings.
  • src/hydroua_gpt/preprocess/ — PDF extraction and corpus creation.
  • src/hydroua_gpt/generation/ — vLLM-powered QA generation.
  • src/hydroua_gpt/finetune/ — dataset preparation and LoRA/QLoRA training.
  • src/hydroua_gpt/inference/ — model loading and hydrology chat inference.
  • src/hydroua_gpt/rag/ — literature collection, embeddings, FAISS indexing, and retrieval helpers.
  • src/hydroua_gpt/app/ — Gradio application interface.
  • scripts/slurm/ — example GPU/HPC job scripts for QA generation, fine-tuning, and app deployment.

Technical Stack

  • LLM frameworks: PyTorch, Transformers, PEFT, Accelerate
  • Fine-tuning: LoRA, optional QLoRA, BF16, gradient checkpointing, gradient accumulation
  • Large-scale generation: vLLM with tensor parallelism
  • Data processing: JSONL pipelines, PyMuPDF, YAML configs, reproducible CLI scripts
  • RAG: SentenceTransformers, FAISS, OpenAlex literature collection
  • Application: Gradio hydrology assistant interface
  • Deployment context: Linux GPU/HPC environment with SLURM job examples

Impact

HydroUA-GPT demonstrates how domain-specific scientific LLMs can be engineered beyond a simple prompt wrapper. The project provides a reusable foundation for adapting open LLMs to hydrology, preserving scientific context, training with parameter-efficient adapters, grounding answers with retrieval, and deploying an interactive assistant for hydrology education and research support.

Responsible AI and Data Handling

The repository is designed to keep copyrighted PDFs, generated derivative datasets, model weights, LoRA adapters, checkpoints, and vector databases out of version control. This keeps the public codebase clean while allowing the training pipeline to be reproduced in a controlled research environment with properly licensed source materials.