codefinetuner 0.1.0__py3-none-any.whl

codefinetuner/finetune/run.py

@@ -0,0 +1,136 @@
+import argparse
+import logging.config
+import shutil
+import sys
+from typing import Tuple
+from pathlib import Path
+
+from datasets import Features, IterableDataset, Sequence, Value, load_dataset
+from transformers import AutoTokenizer
+
+from .config import Config
+from .model import load_and_configure_lora_model
+from .train import plot_loss, save_log, train_lora_model, merge_lora_and_save
+
+
+logger = logging.getLogger(__name__)
+
+
+def _setup_logger(log_level: str) -> None:
+    """
+    Configure the root logger ("") to capture logs from the entry point (__main__),
+    all internal sub-modules, and third-party libraries via a single handler.
+    """
+    logger_config = {
+        "version": 1,
+        "disable_existing_loggers": False,
+        "formatters": {
+            "standard": {
+                "format": "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+            }
+        },
+        "handlers": {
+            "stderr_handler": {
+                "level": log_level,
+                "class": "logging.StreamHandler",
+                "formatter": "standard",
+            }
+        },
+        "loggers": {
+            "": {  # "" corresponds to root logger 
+                "handlers": ["stderr_handler"],
+                "level": log_level,
+                "propagate": False,
+            }
+        }
+    }
+    logging.config.dictConfig(logger_config)
+
+
+def _parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Start or resume LoRA model training")
+    parser.add_argument(
+        "--config",
+        type=Path,
+        required=True,
+        help="YAML config file path with 'finetune:' section. All Config.MISSING fields must be provided."
+    )
+    parser.add_argument(
+        "--log-level",
+        type=str,
+        default="INFO",
+        choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
+        help="Logging level.",
+    )
+    args = parser.parse_args()
+    return args
+
+
+def _ensure_clean_checkpoint_dir(config: Config) -> None:
+    checkpoints_dir = config.trainer_checkpoints_dir_path
+    checkpoint = config.trainer_resume_from_checkpoint
+    clear = config.trainer_clear_checkpoint_dir
+    if checkpoint and not clear:
+        logger.info(f"Resuming from: {checkpoint}")
+    elif not checkpoint and clear: 
+        if checkpoints_dir.exists():
+            logger.warning(f"Deleting: {checkpoints_dir}")
+            shutil.rmtree(checkpoints_dir)
+        checkpoints_dir.mkdir(parents=True, exist_ok=True)
+        logger.info(f"Cleared: {checkpoints_dir}")
+    elif checkpoint and clear: 
+        raise RuntimeError("Configuration conflict: Cannot resume from a checkpoint while '--delete-all-checkpoints' is set.")
+    else:
+        logger.info(f"Starting fresh training. Existing checkpoints in {checkpoints_dir} are preserved.")
+
+
+def load_datasets(config: Config) -> Tuple[IterableDataset, IterableDataset]:
+    # Define the expected schema/features of datasets.
+    # Use 'int32' which the datasets library and pytorch map correctly to int tensors.
+    dataset_features = Features({
+        'input_ids': Sequence(feature=Value(dtype='int32')),
+        'attention_mask': Sequence(feature=Value(dtype='int32')),
+        'labels': Sequence(feature=Value(dtype='int32')),
+    })
+
+    # Enable streaming mode to load the dataset as an iterator.
+    # This allows processing data samples on-the-fly without downloading or loading the entire dataset into memory.
+    # https://huggingface.co/docs/datasets/stream
+    train_dataset = load_dataset("json", data_files=str(config.train_dataset_path), features=dataset_features, streaming=True)["train"]
+    train_dataset = train_dataset.shuffle(buffer_size=config.dataset_shuffle_buffer_size, seed=config.dataset_shuffle_seed)  # take up to suffle_buffer_size examples and randomly shuffle them
+    eval_dataset = load_dataset("json", data_files=str(config.eval_dataset_path), features=dataset_features, streaming=True)["train"]
+    return train_dataset, eval_dataset 
+
+
+def run(config: Config) -> None:
+    _ensure_clean_checkpoint_dir(config)
+
+    train_dataset, eval_dataset= load_datasets(config)
+    logger.info(f"Dataset: {config.dataset_train_dataset_length} train examples, max_steps={config.trainer_max_steps}")
+
+    lora_model = load_and_configure_lora_model(config)
+    lora_model.print_trainable_parameters()
+
+    tokenizer = AutoTokenizer.from_pretrained(config.model_name)
+    tokenizer.pad_token = config.fim_pad_token
+    tokenizer.padding_side = "right"
+
+    log_history = train_lora_model(config, lora_model, tokenizer, train_dataset, eval_dataset)
+    merge_lora_and_save(config, tokenizer)
+    save_log(config, log_history)
+
+    plot_loss(config)
+
+
+def main() -> None:
+    user_args = _parse_args()
+    _setup_logger(user_args.log_level)
+    try:
+        finetune_config = Config.load_from_yaml(user_args.config)
+        run(finetune_config)
+    except Exception as e:
+        logger.exception(f"Finetuning failed")
+        sys.exit(1)
+
+if __name__ == "__main__":
+    main()

codefinetuner/finetune/train.py

@@ -0,0 +1,204 @@
+import gc
+import json
+import logging
+import shutil
+from pathlib import Path
+from typing import Dict, List
+
+import matplotlib.pyplot as plt
+import torch
+from datasets import IterableDataset
+from peft import PeftModel
+from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
+
+from .config import Config
+
+
+logger = logging.getLogger(__name__)
+
+
+class FIMDataCollator:
+    def __init__(self, tokenizer, label_pad_token_id):
+        self.tokenizer = tokenizer
+        self.label_pad_token_id = label_pad_token_id 
+
+    def __call__(self, examples: List[Dict[str, List[int]]]) -> Dict[str, torch.Tensor]:
+        # find max length for this specific batch
+        max_ex_length = 0 
+        for ex in examples:
+            if len(ex["input_ids"]) > max_ex_length:
+                max_ex_length = len(ex["input_ids"])
+        
+        # apply padding 
+        padded_examples = [] 
+        for ex in examples:
+            pad_length = max_ex_length - len(ex["input_ids"])
+            padded_ex = {  
+                "input_ids": (ex["input_ids"] + [self.tokenizer.pad_token_id] * pad_length),  
+                "attention_mask": (ex["attention_mask"] + [0] * pad_length),  
+                "labels": (ex["labels"] + [self.label_pad_token_id] * pad_length)  
+            }  
+            padded_examples.append(padded_ex)
+
+        # stack separate examples into a single tensors matrix
+        examples_batch = {
+            "input_ids": torch.tensor([ex["input_ids"] for ex in padded_examples], dtype=torch.long),
+            "attention_mask": torch.tensor([ex["attention_mask"] for ex in padded_examples], dtype=torch.long),
+            "labels": torch.tensor([ex["labels"] for ex in padded_examples], dtype=torch.long)
+        }
+
+        return examples_batch
+    
+
+def train_lora_model(
+        config: Config,
+        lora_model: PeftModel,
+        tokenizer: AutoTokenizer,
+        train_dataset: IterableDataset,
+        eval_dataset: IterableDataset,
+) -> List: 
+    logger.info(f"Starting training: {config.trainer_max_steps} steps on {config.device}, batch_size={config.trainer_per_device_train_batch_size}") 
+
+    if config.model_dtype == torch.bfloat16:
+        trainer_bf16 = True
+        trainer_fp16 = False
+    elif config.model_dtype == torch.float16:
+        trainer_bf16 = False 
+        trainer_fp16 = True 
+    else:
+        trainer_bf16 = False 
+        trainer_fp16 = False 
+
+    training_args = TrainingArguments(
+        output_dir=config.trainer_checkpoints_dir_path,
+        per_device_train_batch_size=config.trainer_per_device_train_batch_size,
+        per_device_eval_batch_size=config.trainer_per_device_eval_batch_size,
+        gradient_accumulation_steps=config.trainer_gradient_accumulation_steps, 
+        learning_rate=config.trainer_learning_rate,
+        weight_decay=config.trainer_weight_decay,
+        max_grad_norm=config.trainer_max_grad_norm,
+        max_steps=config.trainer_max_steps,
+        lr_scheduler_type=config.trainer_lr_scheduler_type,
+        warmup_steps=config.trainer_warmup_steps,
+        eval_strategy=config.trainer_eval_strategy,
+        logging_steps=config.trainer_logging_steps,
+        eval_steps=config.trainer_eval_steps,
+        logging_strategy=config.trainer_logging_strategy,
+        save_strategy=config.trainer_save_strategy,
+        save_steps=config.trainer_save_steps,
+        bf16=trainer_bf16,  
+        fp16=trainer_fp16,  
+        gradient_checkpointing=config.trainer_gradient_checkpointing
+    )
+
+    data_collator = FIMDataCollator(
+        tokenizer=tokenizer,
+        label_pad_token_id=config.label_pad_token_id
+    )
+    
+    trainer = Trainer(
+        model=lora_model,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        processing_class = tokenizer,
+        data_collator = data_collator
+    )
+
+    checkpoint = config.trainer_resume_from_checkpoint
+    if checkpoint == "last":
+        trainer.train(resume_from_checkpoint=True)
+    elif checkpoint is not None:
+        trainer.train(resume_from_checkpoint=checkpoint) 
+    else:
+        trainer.train()  # train from scratch
+
+    lora_model.save_pretrained(config.lora_adapter_path) # save lora adapter only
+    log_history = trainer.state.log_history
+
+    del trainer
+    del lora_model
+    return log_history
+
+
+def merge_lora_and_save(config: Config, tokenizer: AutoTokenizer) -> None:
+    gc.collect()  # force garbage collection
+
+    if torch.cuda.is_available():
+        torch.cuda.synchronize()  # wait for gpu
+        torch.cuda.empty_cache()  # clear gpu cache
+    
+    # ensure offload folder for merging exists before loading the model
+    config.trainer_model_merge_offload_folder_path.mkdir(parents=True, exist_ok=True)
+
+    # load fresh base model
+    base_model = AutoModelForCausalLM.from_pretrained(
+        pretrained_model_name_or_path=config.model_name,
+        dtype=config.model_dtype,
+        device_map="auto",
+        offload_folder=str(config.trainer_model_merge_offload_folder_path),  # offload model layers to disk during loading to prevent RAM (OOM) crashes 
+        low_cpu_mem_usage=True
+    )
+
+    lora_model = PeftModel.from_pretrained(
+        model=base_model, 
+        model_id=config.lora_adapter_path,
+        offload_folder=str(config.trainer_model_merge_offload_folder_path)
+    )
+
+    # merge lora adapter into base model and save it with the tokenizer of the model
+    merged_model= lora_model.merge_and_unload() 
+    merged_model = merged_model.to(config.model_dtype) 
+    merged_model.save_pretrained(config.lora_model_path)
+    tokenizer.save_pretrained(config.lora_model_path)
+
+    # clean up offload folder
+    if config.trainer_model_merge_offload_folder_path.exists():
+        shutil.rmtree(config.trainer_model_merge_offload_folder_path)
+
+
+def save_log(config: Config, log_history: List) -> None:
+    history = {
+        "train": {"steps": [], "loss": [], "learning_rate": [], "epoch": []},
+        "eval": {"steps": [], "loss": [], "epoch": []}
+    }
+
+    for entry in log_history:
+        # Training logs
+        if "loss" in entry:
+            history["train"]["loss"].append(entry["loss"])
+            history["train"]["steps"].append(entry["step"])
+            history["train"]["epoch"].append(entry.get("epoch"))
+            history["train"]["learning_rate"].append(entry.get("learning_rate"))
+        
+        # Evaluation logs
+        elif "eval_loss" in entry:
+            history["eval"]["loss"].append(entry["eval_loss"])
+            history["eval"]["steps"].append(entry["step"])
+            history["eval"]["epoch"].append(entry.get("epoch"))
+
+    log_path = config.trainer_log_path 
+    with log_path.open("w", encoding="utf-8") as f:
+        json.dump(history, f, indent=2)
+
+
+def plot_loss(config: Config) -> None:
+    log_path = config.trainer_log_path 
+    with log_path.open("r", encoding="utf-8") as f:
+        data = json.load(f)
+
+    plt.figure(figsize=(8, 5))
+    
+    if data["train"]["steps"]:
+        plt.plot(data["train"]["steps"], data["train"]["loss"], label="Train Loss")
+    
+    if data["eval"]["steps"]:
+        plt.plot(data["eval"]["steps"], data["eval"]["loss"], label="Eval Loss", marker='o')
+
+    plt.xlabel("Step")
+    plt.ylabel("Loss")
+    plt.title("Training and Evaluation Loss")
+    plt.legend()
+    plt.grid(True)
+    plt.savefig(config.trainer_plot_path)
+    plt.close()

codefinetuner/pipeline.py

@@ -0,0 +1,152 @@
+import argparse
+import logging
+import logging.config
+import sys
+from pathlib import Path
+
+from .preprocess.config import Config as PreprocessConfig
+from .preprocess.run import run as preprocess_run
+from .finetune.config import Config as FinetuneConfig
+from .finetune.run import run as finetune_run
+from .evaluate.config import Config as EvaluateConfig
+from .evaluate.run import run as evaluate_run
+from .convert.config import Config as ConvertConfig
+from .convert.run import run as convert_run
+
+
+logger = logging.getLogger(__name__)
+
+
+def _setup_logger(log_level: str) -> None:
+    """
+    Configure the root logger ("") to capture logs from the entry point (__main__),
+    all internal sub-modules, and third-party libraries via a single handler.
+    """
+    logger_config = {
+        "version": 1,
+        "disable_existing_loggers": False,
+        "formatters": {
+            "standard": {
+                "format": "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+            }
+        },
+        "handlers": {
+            "stderr_handler": {
+                "level": log_level,
+                "class": "logging.StreamHandler",
+                "formatter": "standard",
+            }
+        },
+        "loggers": {
+            "": {  # "" corresponds to root logger
+                "handlers": ["stderr_handler"],
+                "level": log_level,
+                "propagate": False,
+            }
+        },
+    }
+    logging.config.dictConfig(logger_config)
+
+
+def _parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Run full pipeline: preprocess -> finetune -> evaluate."
+    )
+    parser.add_argument(
+        "--config",
+        type=Path,
+        required=True,
+        help="YAML config file path with 'preprocess:', 'finetune:' and 'evaluate:' sections.",
+    )
+    parser.add_argument(
+        "--log-level",
+        type=str,
+        default="INFO",
+        choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
+        help="Logging level for the pipeline entrypoint.",
+    )
+    parser.add_argument(
+        "--skip-preprocess",
+        action="store_true",
+        help="Skip preprocess stage and reuse existing datasets.",
+    )
+    parser.add_argument(
+        "--skip-finetune",
+        action="store_true",
+        help="Skip finetune stage and reuse existing checkpoints.",
+    )
+    parser.add_argument(
+        "--skip-evaluate",
+        action="store_true",
+        help="Skip evaluation stage.",
+    )
+    parser.add_argument(
+        "--skip-convert",
+        action="store_true",
+        help="Skip conversion to gguf file stage.",
+    )
+    return parser.parse_args()
+
+
+def run_pipeline(
+    config_path: Path,
+    skip_preprocess: bool = False,
+    skip_finetune: bool = False,
+    skip_evaluate: bool = False,
+    skip_convert: bool = False,
+) -> None:
+    logger.info("Starting end-to-end pipeline")
+
+    if not skip_preprocess:
+        logger.info("=== Stage 1/4: Preprocess ===")
+        preprocess_config = PreprocessConfig.load_from_yaml(config_path)
+        preprocess_run(preprocess_config)
+        logger.info("Finished preprocess stage")
+    else:
+        logger.info("Skipping preprocess stage")
+
+    if not skip_finetune:
+        logger.info("=== Stage 2/4: Finetune ===")
+        finetune_config = FinetuneConfig.load_from_yaml(config_path) 
+        finetune_run(finetune_config)
+        logger.info("Finished finetune stage")
+    else:
+        logger.info("Skipping finetune stage")
+
+    if not skip_evaluate:
+        logger.info("=== Stage 3/4: Evaluate ===")
+        evaluate_config = EvaluateConfig.load_from_yaml(config_path)
+        evaluate_run(evaluate_config)
+        logger.info("Finished evaluate stage")
+    else:
+        logger.info("Skipping evaluate stage")
+    
+    if not skip_convert:
+        logger.info("=== Stage 4/4: Convert ===")
+        convert_config = ConvertConfig.load_from_yaml(config_path)
+        convert_run(convert_config)
+        logger.info("Finished conversion stage")
+    else:
+        logger.info("Skipping conversion stage")
+
+    logger.info("Pipeline completed successfully")
+
+
+def main() -> None:
+    user_args = _parse_args()
+    _setup_logger(user_args.log_level)
+    try:
+        run_pipeline(
+            config_path=user_args.config,
+            skip_preprocess=user_args.skip_preprocess,
+            skip_finetune=user_args.skip_finetune,
+            skip_evaluate=user_args.skip_evaluate,
+            skip_convert=user_args.skip_convert,
+        )
+    except Exception:
+        logger.exception("Pipeline execution failed")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

codefinetuner/preprocess/config.py

@@ -0,0 +1,160 @@
+import logging
+import math
+import json
+from importlib import resources
+from dataclasses import dataclass, field, fields
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+from omegaconf import OmegaConf, MISSING
+
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class Config:
+    # --- Mandatory Parameters ---
+    model_name: str = MISSING
+    fim_prefix_token: str = MISSING
+    fim_middle_token: str = MISSING
+    fim_suffix_token: str = MISSING
+    fim_pad_token: str = MISSING
+    eos_token: str = MISSING
+    data_language: str = MISSING
+    data_extensions: list[str] = MISSING
+
+    # --- Preprocess Local Parameters ---
+    split_mode: str = "auto"
+    train_ratio: float = 0.8
+    eval_ratio: float = 0.1
+    test_ratio: float = 0.1
+    max_token_sequence_length: int = 1024  # used with bytes_per_token_ratio to convert bytes to tokens, final token count is thus not exact
+    max_code_blocks_ast_depth: int = 2  # depth 1 is root, 2 includes child nodes (e.g. functions)
+    min_middle_tokens_length: int = 20  # used with estimated bytes_per_token_ratio to convert bytes to tokens, final token count is thus not exact 
+    max_middle_tokens_length: int = 200  # used with estimated bytes_per_token_ratio to convert bytes to tokens, final token count is thus not exact
+    fim_examples_per_subblock_ratio: float = 1.0  # 1.0 = all fim examples of a subblock are extracted, 0.5 = onls 50% of fim examples of a subblock are extracted
+    tokenizer_batch_size: int = 32
+    rng_seed: int = 0 
+
+    # --- Tree Sitter Parser ---
+    tree_sitter_parser: Any = field(init=False)  # type hint Any because omegaconf does not recognize ts.Parser as a valid type (yaml file cannot contain an object of ts.Parser)
+    tree_sitter_block_types: Any = field(init=False)  # type hint Any because omegaconf does not support set type
+    tree_sitter_subblock_types: Any = field(init=False)  # type hint Any because omegaconf does not support set type
+
+    # --- Paths ---
+    workspace_path: Path | None = None 
+    raw_data_path: Path | None = None 
+    train_dataset_path: Path = field(init=False)
+    eval_dataset_path: Path = field(init=False)
+    test_dataset_path: Path = field(init=False)
+    tree_sitter_parser_path: Path | None = None
+    tree_sitter_definitions_path: Path | None = None
+
+    # --- Randomization ---
+    rng: Any = field(init=False)  # type hint Any because omegaconf does not support np.random.Generator type
+
+    @classmethod
+    def load_from_yaml(cls, yaml_path: Path) -> "Config":
+        if not yaml_path.exists():
+            raise FileNotFoundError(f"Configuration file not found at: {yaml_path}")
+        
+        config_dict = OmegaConf.structured(cls)
+        try:
+            yaml_file_node = OmegaConf.load(yaml_path)
+        except Exception as e:
+            raise ValueError(f"Failed to parse YAML file: {yaml_path}") from e
+            
+        yaml_file_dict = OmegaConf.to_container(yaml_file_node, resolve=True)
+        yaml_preprocess_dict = yaml_file_dict.get("preprocess", {})
+
+        yaml_preprocess_valid_dict = {}
+        # Filter YAML fields to include only those defined in the Config dataclass.
+        # This prevents OmegaConf from raising an AttributeError when encountering 
+        # global YAML anchors or keys not present in the current Config dataclass. 
+        for field in fields(cls):
+            if field.name in yaml_preprocess_dict:
+                yaml_preprocess_valid_dict[field.name] = yaml_preprocess_dict[field.name]
+        logger.debug(f"Filtered YAML configuration: {yaml_preprocess_valid_dict}")
+
+        merged_config_dict = OmegaConf.merge(config_dict, yaml_preprocess_valid_dict)
+        
+        return OmegaConf.to_object(merged_config_dict)
+
+    def __post_init__(self) -> None:
+        self._validate_ratio()
+        self._setup_paths()
+        self._ensure_output_paths_exist()
+        self._load_language_blocks()
+        self._init_tree_sitter_parser()
+        self.rng = np.random.default_rng(seed=self.rng_seed)
+        logger.debug("Config initialization complete.")
+    
+    def _validate_ratio(self):
+        total_ratio = self.train_ratio + self.eval_ratio + self.test_ratio
+        if not math.isclose(total_ratio, 1.0, rel_tol=1e-6):
+            raise ValueError(f"Train + eval + test ratios must sum to 1.0, got {total_ratio}")
+
+    def _setup_paths(self) -> None:
+        if self.workspace_path is None:
+            self.workspace_path = Path.cwd()
+        if self.raw_data_path is None:
+            self.raw_data_path = self.workspace_path / "data"
+            
+        if self.tree_sitter_definitions_path is None:
+            try:
+                self.tree_sitter_definitions_path = Path(resources.files(__package__)/ "tree_sitter_definitions.json")
+            except Exception:
+                raise FileNotFoundError("Missing internal tree_sitter_definitions.json") 
+
+        self.preprocess_outputs_dir_path = self.workspace_path / "outputs" / "preprocess"
+        self.train_dataset_path = self.preprocess_outputs_dir_path / "results" / "datasets" / "train_dataset.jsonl"
+        self.eval_dataset_path = self.preprocess_outputs_dir_path / "results" / "datasets" / "eval_dataset.jsonl"
+        self.test_dataset_path = self.preprocess_outputs_dir_path / "results" / "datasets" / "test_dataset.jsonl"
+        logger.debug(f"Resolved workspace path to: {self.workspace_path}")
+
+    def _ensure_output_paths_exist(self) -> None:
+        paths = [
+            self.preprocess_outputs_dir_path,
+            self.train_dataset_path,
+            self.eval_dataset_path,
+            self.test_dataset_path
+        ]
+
+        for path in paths:
+            if not path.parent.exists():
+                path.parent.mkdir(parents=True, exist_ok=True)
+                logger.debug(f"Created parent directory: {path.parent}")
+            else:
+                logger.debug(f"Parent directory already exists: {path.parent}")  
+ 
+    def _load_language_blocks(self) -> None:
+        blocks_path = self.tree_sitter_definitions_path 
+        with open(blocks_path, "r", encoding="utf-8") as f:
+            language_data = json.load(f)
+
+        language_blocks = language_data.get(self.data_language)
+        if language_blocks is None:
+            raise ValueError(f"Language '{self.data_language}' not found in {blocks_path}")
+
+        tree_sitter_block_types = language_blocks.get("block_types")
+        tree_sitter_subblock_types = language_blocks.get("subblock_types")
+        if not isinstance(tree_sitter_block_types, list) or not isinstance(tree_sitter_subblock_types, list):
+            raise ValueError(f"Invalid block definitions for '{self.data_language}' in {blocks_path}")
+
+        self.tree_sitter_block_types = set(tree_sitter_block_types)
+        self.tree_sitter_subblock_types = set(tree_sitter_subblock_types)
+
+    def _init_tree_sitter_parser(self) -> None:
+        if self.tree_sitter_parser_path:
+            logger.info(f"Loading custom tree-sitter parser from {self.tree_sitter_parser_path}")
+            from .extract import get_custom_tree_sitter_parser
+            self.tree_sitter_parser = get_custom_tree_sitter_parser(self.tree_sitter_parser_path, self.data_language)
+        else:
+            logger.info(f"Loading tree-sitter language pack parser {self.data_language}")
+            from .extract import get_tree_sitter_language_pack_parser
+            self.tree_sitter_parser = get_tree_sitter_language_pack_parser(self.data_language)
+        
+        if self.tree_sitter_parser is None:
+            raise RuntimeError("Tree-sitter parser not initialized")