feat(benchmark): benchmark suite + register_models export

.gitignore

@@ -6,4 +6,11 @@
 /docs/build/
 .env
 .env.example
-JuliaHealthLLM_exp_raw_combined.txt
+JuliaHealthLLM_exp_raw_combined.txt
+
+# benchmark/data files
+benchmark_result_*.json
+synthea_1M_3YR.duckdb
+train.jsonl
+nul
+FunSQLQueries/

Project.toml

@@ -5,6 +5,7 @@ authors = ["ParamThakkar123 <paramthakkar864@gmail.com> and TheCedarPrince <jaco
 
 [deps]
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
+Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 DrWatson = "634d3b9d-ee7a-5ddf-bec9-22491ea816e1"
 DuckDB = "d2f5444f-75bc-4fdf-ac35-56f514c445e1"
 FunSQL = "cf6cc811-59f4-4a10-b258-a8547a8f6407"
@@ -13,14 +14,17 @@ HuggingFaceHub = "d0076355-e2c0-48e6-a044-05906e51b7fc"
 JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
 LibPQ = "194296ae-ab2e-5f79-8cd4-7183a0a5a0d1"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
 PromptingTools = "670122d1-24a8-4d70-bfce-740807c42192"
 RAGTools = "16ddad29-bbe8-45a7-857d-3d9514eb0023"
 Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
 
 [compat]
 DataFrames = "1.8.0"
+Dates = "1.10"
 DrWatson = "2.19.1"
 DuckDB = "1.3.2"
 FunSQL = "0.15.0"
@@ -29,11 +33,13 @@ HuggingFaceHub = "0.1.2"
 JSON3 = "1.14.3"
 LibPQ = "1.18.0"
 LinearAlgebra = "1.10"
+ProgressMeter = "1.11.0"
 PromptingTools = "0.82.1"
 RAGTools = "0.7.0"
 Serialization = "1.10"
 SparseArrays = "1.10"
 Statistics = "1.10"
+Tables = "1.12.1"
 julia = "1.10"
 
 [extras]

docs/src/index.md

@@ -10,5 +10,5 @@ Documentation for [HealthLLM](https://github.com/ParamThakkar123/HealthLLM.jl).
 ```
 
 ```@autodocs
-Modules = [HealthLLM, HealthLLM.Utils]
+Modules = [HealthLLM, HealthLLM.Utils, HealthLLM.Benchmark]
 ```

scripts/run_benchmark_llama.jl

@@ -0,0 +1,250 @@
+#!/usr/bin/env julia
+# Run the FunSQL benchmark using HealthLLM against the FunSQL dataset.
+# Usage:
+#   julia scripts/run_benchmark_llama.jl [model_name] [model_embedding] [sample_limit]
+# Example:
+#   julia scripts/run_benchmark_llama.jl "meta-llama/Llama-2-7b-chat-hf" "sentence-transformers/all-MiniLM-L6-v2" 20
+
+using HealthLLM
+using HuggingFaceHub
+const HF = HuggingFaceHub
+using JSON3
+using Dates
+# progress bar (required)
+using ProgressMeter
+using HTTP
+const HAS_PROGRESS = true
+
+function download_datasets(; synthea_path::Union{Nothing,String}=nothing, funsql_path::Union{Nothing,String}=nothing)
+    # Allow using local copies of datasets if present on disk
+    local_default_synthea = "E:\\HealthLLM.jl\\synthea_1M_3YR.duckdb"
+    if synthea_path === nothing && isfile(local_default_synthea)
+        println("Found local Synthea duckdb at: $local_default_synthea; using it.")
+        synthea_path = local_default_synthea
+    end
+
+    # Prefer a locally checked-out FunSQL dataset when available
+    local_default_funsql = "E:\\HealthLLM.jl\\FunSQLQueries\\train.jsonl"
+    if funsql_path === nothing && isfile(local_default_funsql)
+        println("Found local FunSQL dataset at: $local_default_funsql; using it.")
+        funsql_path = local_default_funsql
+    end
+
+    # If caller provided both paths, prefer those
+    if synthea_path !== nothing && funsql_path !== nothing
+        println("Using provided dataset paths.")
+        return synthea_path, funsql_path
+    end
+
+    println("Downloading datasets from HuggingFace (or using local cache)...")
+    synthea_name = "JuliaHealthOrg/JuliaHealthDatasets"
+    funsql_name = "JuliaHealthOrg/FunSQLQueries"
+    synthea_ds = HF.info(HF.Dataset, synthea_name)
+    funsql_ds = HF.info(HF.Dataset, funsql_name)
+
+    # show which files will be downloaded and where they will be stored
+    println("Will download (only missing files):")
+    println(" - $synthea_name -> synthea_1M_3YR.duckdb")
+    println(" - $funsql_name -> train.jsonl")
+
+    # helper: attempt a streaming HTTP download with per-file progress
+    function try_stream_download(repo::String, filename::String)
+        candidates = [
+            "https://huggingface.co/datasets/$repo/resolve/main/$filename",
+            "https://huggingface.co/$repo/resolve/main/$filename",
+            "https://huggingface.co/datasets/$repo/resolve/refs/heads/main/$filename"
+        ]
+
+        for url in candidates
+            try
+                # Try HEAD first to learn content length
+                head_res = try
+                    HTTP.request("HEAD", url)
+                catch
+                    nothing
+                end
+
+                total_bytes = nothing
+                if head_res !== nothing && head_res.status == 200
+                    clen = get(head_res.headers, "Content-Length", nothing)
+                    if clen !== nothing
+                        try
+                            total_bytes = parse(Int, String(clen))
+                        catch
+                            total_bytes = nothing
+                        end
+                    end
+                end
+
+                dest = abspath(filename)
+                println("Attempting download: $filename from $url -> $dest")
+
+                # Create a Progress instance in a way compatible with multiple
+                # ProgressMeter versions. Older versions may not accept the
+                # `show_eta` keyword, so try the full call first and fall back
+                # to a minimal constructor when necessary.
+                function _make_progress(n)
+                    try
+                        return Progress(n; show_eta=true)
+                    catch
+                        return Progress(n)
+                    end
+                end
+
+                pm_file = total_bytes !== nothing ? _make_progress(total_bytes) : _make_progress(1)
+
+                # Stream GET
+                try
+                    HTTP.open(:GET, url) do stream_io
+                        open(dest, "w") do out_io
+                            bytes_written = 0
+                            while !eof(stream_io)
+                                chunk = read(stream_io, 65536)
+                                if isempty(chunk)
+                                    break
+                                end
+                                write(out_io, chunk)
+                                bytes_written += length(chunk)
+                                try
+                                    if total_bytes !== nothing
+                                        ProgressMeter.update!(pm_file, bytes_written)
+                                    else
+                                        ProgressMeter.update!(pm_file)
+                                    end
+                                catch
+                                    # ignore progress update failures
+                                end
+                            end
+                        end
+                    end
+                catch err
+                    @warn "GET streaming failed for $url: $err"
+                    continue
+                end
+
+                println("Finished downloading $filename to $dest")
+                return dest
+            catch err
+                @warn "Stream download attempt failed for $url: $err"
+            end
+        end
+
+        # fallback to HuggingFaceHub.file_download if available
+        try
+            if isdefined(HF, :file_download)
+                info = HF.info(HF.Dataset, repo)
+                dest = HF.file_download(info, filename)
+                println("Downloaded $filename via HuggingFaceHub.file_download -> $dest")
+                return dest
+            end
+        catch err
+            @warn "HuggingFaceHub.file_download fallback failed: $err"
+        end
+
+        error("Could not download $filename from repository $repo")
+    end
+
+    # overall download progress across both dataset files
+    # Use the same compatibility helper to construct the meter.
+    function _make_progress(n)
+        try
+            return Progress(n; show_eta=true)
+        catch
+            return Progress(n)
+        end
+    end
+
+    # Determine how many files we actually need to download
+    need_synthea = synthea_path === nothing
+    need_funsql = funsql_path === nothing
+    total_to_download = (need_synthea ? 1 : 0) + (need_funsql ? 1 : 0)
+    pm = _make_progress(max(total_to_download, 1))
+
+    synthea_dataset_path = synthea_path
+    if need_synthea
+        synthea_dataset_path = try_stream_download(synthea_name, "synthea_1M_3YR.duckdb")
+        ProgressMeter.update!(pm, 1)
+    end
+
+    funsql_dataset_path = funsql_path
+    if need_funsql
+        # if we already downloaded synthea, update progress position for the second file
+        offset = need_synthea ? 2 : 1
+        funsql_dataset_path = try_stream_download(funsql_name, "train.jsonl")
+        ProgressMeter.update!(pm, offset)
+    end
+
+    return synthea_dataset_path, funsql_dataset_path
+end
+
+function main()
+    model_name = length(ARGS) >= 1 ? ARGS[1] : "meta-llama/Llama-2-7b-chat-hf"
+    model_embedding = length(ARGS) >= 2 ? ARGS[2] : "sentence-transformers/all-MiniLM-L6-v2"
+    sample_limit = length(ARGS) >= 3 ? parse(Int, ARGS[3]) : 10
+
+    println("Model: $model_name")
+    println("Embedding model: $model_embedding")
+    println("Sample limit: $sample_limit")
+
+    # Allow overriding dataset paths via ARGS (positions 4 and 5)
+    provided_synthea = length(ARGS) >= 4 ? ARGS[4] : nothing
+    provided_funsql = length(ARGS) >= 5 ? ARGS[5] : nothing
+
+    synthea_path, funsql_path = download_datasets(synthea_path=provided_synthea, funsql_path=provided_funsql)
+
+    println("Registering models with PromptingTools...")
+    HealthLLM.register_models(model_name, model_embedding)
+
+    # detect number of examples in FunSQL dataset to set progress total
+    detected_len = 0
+    open(funsql_path, "r") do io
+        for line in eachline(io)
+            if !isempty(strip(line))
+                detected_len += 1
+            end
+        end
+    end
+    # if file looks like a single-line JSON array, parse to get length
+    if detected_len <= 1
+        arr = JSON3.read(read(funsql_path, String))
+        detected_len = length(arr)
+    end
+
+    total_examples = sample_limit > 0 ? min(sample_limit, detected_len) : detected_len
+
+    # create ProgressMeter with exact total and register a callback to update it
+    # Use a compatibility constructor to avoid passing unsupported keywords.
+    function _make_progress(n)
+        try
+            return Progress(n; show_eta=true)
+        catch
+            return Progress(n)
+        end
+    end
+
+    pm = _make_progress(total_examples)
+    HealthLLM.register_progress!((current,total,msg)->begin
+        # set the meter to current (clamp to total)
+        v = clamp(current, 0, total_examples)
+        ProgressMeter.update!(pm, v)
+    end)
+
+    println("Starting benchmark run...")
+    res = HealthLLM.run_benchmark(model_name, model_embedding, synthea_path, funsql_path; sample_limit=sample_limit)
+
+    println() # newline after progress
+    println("Metrics:")
+    println(JSON3.write(res["metrics"]))
+
+    failures = res["failures"]
+    println("Failures: $(length(failures))")
+
+    out_file = "benchmark_result_$(replace(string(now()), ':' => '-')).json"
+    open(out_file, "w") do io
+        JSON3.write(io, res)
+    end
+
+    println("Wrote full results to: $out_file")
+end
+
+main()

src/HealthLLM.jl

@@ -12,13 +12,16 @@ include("pgvector.jl")
 include("database.jl")
 include("embedding.jl")
 include("query.jl")
+include("benchmark.jl")
 
-import .Utils: collect_files_with_extensions, write_combined_file, register_models, load_huggingface_model, HuggingFaceLoadResult
+import .Utils: collect_files_with_extensions, write_combined_file, register_models, register_progress!, clear_progress!, load_huggingface_model, HuggingFaceLoadResult
 import .Embedding: build_index_rag
 import .Database: store_embeddings_pgvector
 import .Query: generate_funsql_query
+import .Benchmark: run_benchmark
 
 export collect_files_with_extensions, write_combined_file, generate_funsql_query,
-build_index_rag, store_embeddings_pgvector, register_models, load_huggingface_model, HuggingFaceLoadResult
+       build_index_rag, store_embeddings_pgvector, register_models, register_progress!, clear_progress!, load_huggingface_model, HuggingFaceLoadResult
+export run_benchmark
 
 end