pytorch_geometric icon indicating copy to clipboard operation
pytorch_geometric copied to clipboard

Enable local TXT2KG + uncertainty-aware LLM with EDFL/ISR

Open leochlon opened this issue 5 months ago • 1 comments

Fixed import path Corrected torch_geometric.nn.LLM → torch_geometric.llm.models.LLM to match the upstream module layout and avoid ImportError.

Local TXT2KG support Added --use_local_txt2kg flag so users can choose between OpenAI/NIM APIs and local Hugging Face models for knowledge graph extraction. Enables fully offline pipelines, avoids API costs, and improves reproducibility.

Checkpoint validation Properly validates cached triples when switching between local vs. API models. Prevents stale or mismatched KG states from silently propagating into training/inference.

GPU utilization Passes n_gpus parameter to TXT2KG’s internal LLM, ensuring multi-GPU hardware is actually used.

Uncertainty-aware LLM Extended LLM(..., uncertainty_estim=True) to compute EDFL/B2T/ISR metrics before generation:

Returns per-item uncertainty metrics: ISR, Δ̄, B2T, RoH bound, priors. Optionally abstains ("[ABSTAIN]") if ISR < threshold. During training, masks labels for low-ISR items, enabling GNN+LLM fine-tuning with calibrated refusal rather than random guessing tests

RAG_TEST=1 pytest test/llm/models/test_llm.py::test_llm_uncertainty -v

leochlon avatar Oct 03 '25 02:10 leochlon

Codecov Report

:x: Patch coverage is 17.64706% with 112 lines in your changes missing coverage. Please review. :white_check_mark: Project coverage is 84.89%. Comparing base (c211214) to head (acad83c). :warning: Report is 128 commits behind head on master.

Files with missing lines Patch % Lines
torch_geometric/llm/models/llm.py 17.96% 105 Missing :warning:
torch_geometric/llm/models/txt2kg.py 12.50% 7 Missing :warning:

:x: Your patch check has failed because the patch coverage (17.64%) is below the target coverage (80.00%). You can increase the patch coverage or adjust the target coverage.

Additional details and impacted files
@@            Coverage Diff             @@
##           master   #10480      +/-   ##
==========================================
- Coverage   86.11%   84.89%   -1.23%     
==========================================
  Files         496      510      +14     
  Lines       33655    36070    +2415     
==========================================
+ Hits        28981    30620    +1639     
- Misses       4674     5450     +776     

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

:rocket: New features to boost your workflow:
  • :snowflake: Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

codecov[bot] avatar Oct 03 '25 16:10 codecov[bot]