AutoSchemaKG Documentation

repository·main·Indexed 21 days ago

https://github.com/hkust-knowcomp/autoschemakg

AutoSchemaKG (atlas-rag v0.0.5.post1) is a framework for fully autonomous knowledge graph construction. It utilizes LLMs for triple extraction of entities and events as well as schema induction, enabling the creation of high-quality, large-scale knowledge graphs from unstructured text without the need for predefined schemas. The repository includes tools for factuality evaluation via the FELM benchmark and general task evaluation using the Language Model Evaluation Harness.

Tokens
64.6K
Snippets
177
Records
223
Agent score
73%

What's inside AutoSchemaKG

  1. Overview of Atlas-RAG in AutoSchemaKG

    main

    Atlas-RAG

    Atlas-RAG is the core package of the AutoSchemaKG framework. It is designed for autonomous knowledge graph (KG) construction by integrating Knowledge Graph Triple Extraction with dynamic Schema Induction.

    Key Capabilities:

    • Automatic Schema Generation: Unlike traditional pipelines that require predefined schemas, Atlas-RAG uses conceptualization to generate schemas automatically.
    • Unstructured Data Processing: Constructs high-quality KGs from unstructured text sources such as PDFs and Markdown.
    • Zero-shot Inferencing: Enables cross-domain inferencing without manual schema tuning.

    This package implements the research from the paper: 'AutoSchemaKG: Autonomous Knowledge Graph Construction through Dynamic Schema Induction from Web-Scale Corpora'.

  2. Overview of the Language Model Evaluation Harness

    main

    The Language Model Evaluation Harness is a unified framework designed to test generative language models across a wide variety of evaluation tasks. It is the backend used for the Hugging Face Open LLM Leaderboard.

    Key Capabilities:

    • Extensive Benchmarks: Includes over 60 standard academic benchmarks with hundreds of subtasks.
    • Diverse Model Support:
      • Hugging Face transformers (including quantization via GPTQModel and AutoGPTQ).
      • GPT-NeoX and Megatron-DeepSpeed.
      • Fast inference via vLLM.
      • Commercial APIs like OpenAI and TextSynth.
      • Evaluation on adapters (e.g., LoRA) via Hugging Face peft.
    • Reproducibility: Uses publicly available prompts to ensure results are comparable across different research papers.
    • Extensibility: Supports custom prompts and evaluation metrics.
  3. Explore task families and descriptions

    main

    The lm-evaluation-harness supports a wide variety of task families across different languages and domains. For detailed information, including precise meanings, sources, and specific task names, refer to the individual README.md files located in each task's subfolder.

    Common task categories include:

    • Reasoning & Math: arc, arithmetic, asdiv, gsm8k, minerva_math.
    • Language Understanding: glue, mmlu, ceval, cmmlu, hellaswag.
    • Coding: humaneval, mbpp, code_x_glue.
    • Multilingual: belebele, global_mmlu, mgsm, xnli.
    • Commonsense: commonsense_qa, piqa, siqa.
    • Medical/Scientific: medqa, pubmedqa, qasper.
  4. Explore AutoSchemaKG examples and tutorials

    main

    The example/ directory contains several specialized workflows:

    • Full Pipeline: example/atlas_full_pipeline.ipynb shows the complete end-to-end KG construction and RAG implementation.
    • Billion-scale KGs: example/atlas_billion_kg_usage.ipynb provides instructions for hosting and using the pre-constructed ATLAS-wiki, ATLAS-pes2o, and ATLAS-cc graphs.
    • Multi-hop QA: example/atlas_multihopqa.ipynb for evaluating performance on benchmarks like MuSiQue, HotpotQA, and 2WikiMultiHopQA.
    • Multilingual: example/multilingual_processing.md for constructing KGs in languages like Chinese, Japanese, and Korean.
    • Customization: example/example_scripts/custom_extraction/ for using custom prompts and schemas.
    • Production/Scaling: example/example_scripts/parallel_generation/ for large-scale processing and example/example_scripts/neo4j_kg/ for hosting KGs via Neo4j APIs.
  5. Parallel Knowledge Graph Generation Overview

    main

    Parallel generation accelerates Knowledge Graph (KG) construction by dividing the dataset into shards and processing each shard simultaneously using separate LLM instances running on different ports.

    This approach consists of two main stages:

    1. Triple Extraction: Extracting KG triples from raw text in parallel.
    2. Concept Generation: Generating concepts and mappings for the extracted entities in parallel.

    By distributing the workload across multiple LLM servers (e.g., multiple vLLM instances), you can significantly reduce the time required for large-scale datasets.

  6. Understand the Retriever hierarchy in AutoSchemaKG

    main

    All retrieval operations in AutoSchemaKG are built upon a hierarchy of base classes. Depending on whether you need to retrieve graph components or text documents, you will interact with different specialized subclasses:

    • BaseRetriever: The abstract base class for all retrieval operations.
    • BaseEdgeRetriever: A specialized class for retrieving edges (triples) from the Knowledge Graph.
    • BasePassageRetriever: A specialized class for retrieving text passages or documents.

    Choosing the correct retriever depends on whether your goal is to find semantic connections in the graph or to find raw text chunks for context.

  7. Configure task groups using GroupConfig

    main

    Task groups are defined using the GroupConfig object. This allows you to bundle multiple tasks together and define how their metrics should be aggregated for reporting.

    Key fields in GroupConfig include:

    • group: The unique name of the group used for CLI invocation.
    • group_alias: A display name used in output tables.
    • task: A single task name or a list of task names that belong to this group.
    • aggregate_metric_list: A list of configurations to aggregate metrics across the subtasks in the group. If empty, no aggregation occurs.
    • metadata: A dictionary for extra configuration, such as overriding the num_fewshot value displayed in results tables.
    # Example conceptual usage of GroupConfig fields
    GroupConfig(
        group="my_group",
        group_alias="My Custom Group",
        task=["task1", "task2"],
        aggregate_metric_list=[
            {
                "metric": "exact_match",
                "aggregation": "mean",
                "weight_by_size": True
            }
        ],
        metadata={"num_fewshot": 5}
    )
  8. Understand MMLU task variants and groups

    main

    The MMLU (Massive Multitask Language Understanding) benchmark in lm_eval is organized into several groups and subgroups. Understanding which variant to use depends on your evaluation methodology:

    Main Groups

    • mmlu: The original multiple-choice benchmark. Choices are provided in the context, and the model is expected to provide the answer letter (e.g., A, B, C, D) in the continuation.
    • mmlu_continuation: A cloze-style variant. It does not include choices in the context and provides the full text of the answer choice in the continuation.
    • mmlu_generation: A generation variant where the LLM is explicitly asked to generate the correct answer letter.

    Subgroups

    Tasks are further categorized into thematic subgroups:

    • mmlu_stem (STEM subjects)
    • mmlu_humanities (Humanities)
    • mmlu_social_sciences (Social Sciences)
    • mmlu_other (Other subjects)

    Subgroup variants are identified by prefixing the subgroup name to the variant type, for example: mmlu_stem_continuation.

  9. Define few-shot examples in YAML or Python

    main

    You can provide few-shot examples for your task in three ways:

    1. Hardcoded in YAML: Use the fewshot_config key with a sampler and a list of samples.
    2. Python function in utils.py: Implement a list_fewshot_samples() -> list[dict] function in an associated utils.py file.
    3. Default: If neither is provided, the framework defaults to using the train, validation, and test sets in that order.

    Note: Each sample dictionary must contain the fields expected by your prompt rendering logic (e.g., if doc_to_text looks for an input key, your samples must include it).

    # Example 1: Hardcoded in YAML
    fewshot_config:
      sampler: first_n
      samples: [
        {"input": "example 1", "output": "answer 1"},
        {"input": "example 2", "output": "answer 2"}
      ]
    
    # Example 2: Python function (in utils.py)
    def list_fewshot_samples() -> list[dict]:
      return [{"input": "ex 1", "output": "ans 1"}, {"input": "ex 2", "output": "ans 2"}]
  10. Handle mixed-language corpora

    main

    When processing a corpus containing multiple languages, you have two strategies:

    Option 1: Unified Processing

    If your triple_extraction_prompt_path contains keys for all languages present in the metadata (e.g., en, zh-CN, ja), you can process the entire corpus in one pass. run_extraction() will automatically route each document to its corresponding prompt.

    Option 2: Language-based Splitting

    For more control, split the corpus into separate JSON files by language before processing. This allows you to run separate extraction pipelines for each language, which is useful for managing different models or configurations per language.

  11. Configure backend-specific parameters in GenerationConfig

    main

    When using native local models (HuggingFace or vLLM offline pipelines), you can pass backend-specific parameters through the GenerationConfig object.

    Supported parameters by backend:

    • vLLM: min_p, use_beam_search, guided_json, guided_regex
    • HuggingFace: repetition_penalty, truncation, padding

    Note: Ensure vllm, transformers, or torch are installed as required by the specific backend.

    from atlas_rag.llm_generator import GenerationConfig
    
    gen_config = GenerationConfig(
        temperature=0.7,
        # vLLM specific
        min_p=0.05,
        guided_json=my_json_schema,
        # HuggingFace specific
        repetition_penalty=1.1
    )