codefinetuner 0.1.0__py3-none-any.whl

codefinetuner/preprocess/extract.py

@@ -0,0 +1,164 @@
+import ctypes
+import logging
+from pathlib import Path
+from typing import Iterator, Tuple
+
+import tree_sitter as ts
+from tree_sitter_language_pack import get_parser
+
+from .config import Config
+
+
+logger = logging.getLogger(__name__)
+
+
+def get_custom_tree_sitter_parser(tree_sitter_lib_path: Path, data_language: str) -> ts.Parser:
+    if not tree_sitter_lib_path.exists():
+        raise FileNotFoundError(f"Library not found: {tree_sitter_lib_path}")
+    try:
+        lib = ctypes.CDLL(str(tree_sitter_lib_path)) # Load C Dynamic Link Library, makes all the public C functions inside the .dylib file available to be called from the Python script.
+        entry_point_func_name = f"tree_sitter_{data_language}"
+        lang_func = getattr(lib, entry_point_func_name) # Retrieves the entry point function of the loaded lib
+        lang_func.restype = ctypes.c_void_p # Tells ctypes that this function returns a C-style pointer (void *)
+        grammar_rules = ts.Language(lang_func()) # Call lang_func() function and wrap the returned raw C pointer into a tree-sitter Language object
+        return ts.Parser(grammar_rules) 
+    except (AttributeError, OSError) as e:
+        raise RuntimeError(f"Failed to load custom tree sitter language parser: {e}") from e
+
+
+def get_tree_sitter_language_pack_parser(data_language: str) -> ts.Parser:
+    return get_parser(data_language)
+
+
+def auto_create_split_paths(config: Config) -> Tuple[list[Path], list[Path], list[Path]]:
+    all_file_paths = []
+    for filepath in config.raw_data_path.rglob("*"):
+        if not filepath.is_file():
+            continue
+        if filepath.suffix.lower() not in config.data_extensions:
+            continue
+        all_file_paths.append(filepath)
+
+    config.rng.shuffle(all_file_paths)
+
+    num_files = len(all_file_paths)
+    if num_files == 0:
+        logger.warning(f"No source files found under {config.raw_data_path} with extensions {config.data_extensions}")
+        return [], [], []
+
+    train_end = int(num_files * config.train_ratio)
+    eval_end = int(num_files * (config.train_ratio + config.eval_ratio))
+
+    train_file_paths = all_file_paths[:train_end]
+    eval_file_paths = all_file_paths[train_end:eval_end]
+    test_file_paths = all_file_paths[eval_end:]
+
+    logger.info(f"Split {num_files} files into {len(train_file_paths)} train, {len(eval_file_paths)} eval, and {len(test_file_paths)} test files.")
+
+    return train_file_paths, eval_file_paths, test_file_paths
+
+
+def _extract_code_blocks_rec(config: Config, node: ts.Node, source_code_utf8: bytes, max_depth: int) -> list[Tuple[bytes, ts.Node]]:
+    """
+    Recursively extract code blocks (e.g. functions) that are used for FIM example generation later 
+    """
+    if max_depth <= 0:
+        return []
+    
+    code_blocks = []
+    if node.type in config.tree_sitter_block_types:
+        code_utf8 = source_code_utf8[node.start_byte:node.end_byte]
+        code_blocks.append((code_utf8, node))
+    
+    for child in node.children:
+        child_blocks = _extract_code_blocks_rec(config, child, source_code_utf8, max_depth-1)
+        code_blocks.extend(child_blocks)  
+    return code_blocks
+
+
+def get_code_blocks_from_paths(config: Config, file_paths: list[Path]) -> Iterator[Tuple[bytes, ts.Node]]:
+    """
+    Generator function yielding code blocks from files one-by-one.
+
+    How generators work:
+    1. Call function: get iterator object (doesn't run code yet)
+    2. next(iterator): processes 1 file, yields first block, then pauses  
+    3. next(iterator): resumes, yields next block, then pauses
+    4. Repeat until end of code blocks, then the generator is exhausted. Generators can only be used once.
+    """
+    for path in file_paths:
+        if not path.is_file():   
+            continue
+        try:
+            source_code_unicode = path.read_text(encoding='utf8')
+        except UnicodeDecodeError:
+            logger.warning(f"Skipping file '{path}': Not a valid UTF-8 file.")
+            continue
+
+        source_code_utf8 = source_code_unicode.encode('utf8')
+        try:
+            tree = config.tree_sitter_parser.parse(source_code_utf8)
+        except Exception as exc:
+            logger.warning(f"Skipping file '{path}': failed to parse with tree-sitter: {exc}")
+            continue
+        root_node = tree.root_node
+
+        code_blocks = _extract_code_blocks_rec(config, root_node, source_code_utf8, max_depth=config.max_code_blocks_ast_depth)
+        config.rng.shuffle(code_blocks)  # shuffle code blocks extracted from a single file
+        for block in code_blocks:
+            yield block  # yields one block at the time 
+
+
+def get_code_blocks_from_auto_split(config: Config) -> Tuple[Iterator[Tuple[bytes, ts.Node]], Iterator[Tuple[bytes, ts.Node]], Iterator[Tuple[bytes, ts.Node]]]:
+    """
+    Auto-split source code files from /data into train/eval/test paths, then extract top-level 
+    code blocks such as functions from each split.
+    """
+    train_file_paths, eval_file_paths, test_file_paths = auto_create_split_paths(config)
+    train_code_blocks_iter = get_code_blocks_from_paths(config, train_file_paths)
+    eval_code_blocks_iter = get_code_blocks_from_paths(config, eval_file_paths)  
+    test_code_blocks_iter =  get_code_blocks_from_paths(config, test_file_paths)
+    return train_code_blocks_iter, eval_code_blocks_iter, test_code_blocks_iter 
+
+
+def _check_required_directories(root_path: Path, required_dirs: list[str]) -> None:
+    missing_paths = []
+    for dir_name in required_dirs:
+        dir_path = root_path / dir_name
+        
+        if not dir_path.is_dir():
+            missing_paths.append(str(dir_path))
+    
+    if missing_paths:
+        raise FileNotFoundError(
+            f"Required directories under {root_path}. "
+            f"Missing directories: {', '.join(missing_paths)}"
+        )
+
+    
+def _get_filtered_paths(config: Config, directory: Path) -> list[Path]:
+    filtered_paths = []
+    
+    for entry in directory.rglob("*"):
+        if not entry.is_file():
+            continue
+        if entry.suffix.lower() not in config.data_extensions:
+            continue
+        filtered_paths.append(entry)
+
+    return filtered_paths
+
+
+def get_code_blocks_from_manual_split(config: Config) -> Tuple[Iterator[Tuple[bytes, ts.Node]], Iterator[Tuple[bytes, ts.Node]], Iterator[Tuple[bytes, ts.Node]]]:
+    """
+    Extract top-level code blocks such as functions from manually split train/eval/test files.
+    """
+    _check_required_directories(config.raw_data_path, ["train", "eval", "test"])
+    train_file_paths = _get_filtered_paths(config, config.raw_data_path / "train")
+    eval_file_paths  = _get_filtered_paths(config, config.raw_data_path / "eval")
+    test_file_paths  = _get_filtered_paths(config, config.raw_data_path / "test")
+
+    train_code_blocks_iter = get_code_blocks_from_paths(config, train_file_paths)
+    eval_code_blocks_iter = get_code_blocks_from_paths(config, eval_file_paths) 
+    test_code_blocks_iter = get_code_blocks_from_paths(config, test_file_paths) 
+    return train_code_blocks_iter, eval_code_blocks_iter, test_code_blocks_iter

codefinetuner/preprocess/process.py

@@ -0,0 +1,180 @@
+import json
+import logging
+from pathlib import Path
+from typing import Iterator, List, Mapping, Tuple
+
+from transformers import AutoTokenizer
+import tree_sitter as ts
+
+from .config import Config
+from .extract import auto_create_split_paths, get_code_blocks_from_paths
+
+
+logger = logging.getLogger(__name__)
+
+
+def estimate_bytes_per_token_ratio(config: Config, tokenizer: AutoTokenizer, number_of_code_blocks: int) -> float:
+    """
+    Estimate bytes per token ratio from first `number_of_code_blocks` code blocks in training split.
+    Always uses auto-split (ignores config.split_mode). 
+    Source code is almost entirely ASCII characters (1 char = 1 byte) 
+    except e.g. specific string literals (printf("π ≈ 3.14159\n");). 
+    ASCII is a subset of UTF-8, so len(bytes) ≈ character count.
+    """
+    train_file_paths, _, _ = auto_create_split_paths(config)  # no matter the split mode always use the auto split
+    total_bytes = 0
+    total_tokens = 0
+    i = 0
+    block_iter = get_code_blocks_from_paths(config, train_file_paths)
+    for block in block_iter:
+        total_bytes += len(block[0])  # bytes from code block
+        tokenized_block = tokenizer(block[0].decode('utf-8')).tokens()
+        total_tokens += len(tokenized_block)
+        i += 1
+        if i >= number_of_code_blocks:
+            break
+    
+    if total_tokens <= 0 or total_bytes <= 0:
+        raise ValueError("Failed to estimate bytes per token ratio")
+    bytes_per_token_ratio = total_bytes / total_tokens 
+    return bytes_per_token_ratio
+
+
+def _extract_subblock_ranges(config: Config, node: ts.Node, base_offset: int) -> list[Tuple[int, int]]:
+    """
+    Recursively depth-first search (DFS) the Abstract syntax tree (AST) and collect all subblock indices.
+    Returns subblock ranges relative to the containing code block.
+    E.g., a subblock range of (2, 30) means the subblock starts at byte position 2
+    and ends at byte position 30, counting from the beginning of the specific code block
+    (not the file).
+    """
+    subblock_ranges = []
+
+    if node.type in config.tree_sitter_subblock_types:
+        relative_start_byte = node.start_byte - base_offset
+        relative_end_byte = node.end_byte - base_offset
+        subblock_ranges.append((relative_start_byte, relative_end_byte))
+    for child in node.children:
+        child_subblock_ranges = _extract_subblock_ranges(config, child, base_offset) 
+        subblock_ranges.extend(child_subblock_ranges)
+        
+    return subblock_ranges
+
+
+def _filter_subblocks(subblock_ranges: list[Tuple[int, int]], max_bytes_per_subblock: int) -> list[Tuple[int, int]]:
+    """
+     Discard subblocks that have a larger end index than max_bytes_per_subblock
+    """
+    subblock_ranges = sorted(subblock_ranges, key=lambda x: x[1]) # sort ranges by end index
+    i = 0
+    while i < len(subblock_ranges) and subblock_ranges[i][1] <= max_bytes_per_subblock:
+        i += 1
+    return subblock_ranges[:i]
+
+
+def _generate_fim_examples_from_code_block(config: Config, code_utf8: bytes, subblock_ranges: list[Tuple[int, int]], bytes_per_token_ratio: float) -> list[bytes]:
+    fim_prefix_token_utf8 = config.fim_prefix_token.encode('utf8')
+    fim_middle_token_utf8 = config.fim_middle_token.encode('utf8')
+    fim_suffix_token_utf8 = config.fim_suffix_token.encode('utf8')
+    eos_token_utf8 = config.eos_token.encode('utf8')
+    
+    num_of_subblocks = len(subblock_ranges)
+    num_of_fim_examples = max(1, int(num_of_subblocks * config.fim_examples_per_subblock_ratio))  
+    unique_random_subblock_indices = config.rng.choice(len(subblock_ranges), size=num_of_fim_examples, replace=False)
+
+    fim_examples = [] 
+
+    for idx in unique_random_subblock_indices:
+        middle_start_byte = subblock_ranges[idx][0]
+        middle_end_byte = subblock_ranges[idx][1]
+        middle_bytes_length = middle_end_byte - middle_start_byte
+        middle_tokens_length = middle_bytes_length / bytes_per_token_ratio
+        # allow only examples within a certain range
+        if (middle_tokens_length < config.min_middle_tokens_length) or (middle_tokens_length > config.max_middle_tokens_length):
+            continue
+
+        prefix = code_utf8[:middle_start_byte]
+        middle = code_utf8[middle_start_byte:middle_end_byte]
+        suffix = code_utf8[middle_end_byte:]
+
+        fim_example = (
+            fim_prefix_token_utf8 + prefix +
+            fim_suffix_token_utf8 + suffix +
+            fim_middle_token_utf8 + middle +
+            eos_token_utf8
+        )
+
+        fim_examples.append(fim_example)
+
+    return fim_examples
+
+
+def create_fim_examples(config: Config, code_blocks_iter: Iterator[Tuple[bytes, ts.Node]], bytes_per_token_ratio: float) -> Iterator[bytes]:
+    """
+    Generator function that takes a generator iterator of code blocks as input, 
+    generates FIM examples from each code block and returns a generator iterator 
+    over all created FIM examples. 
+    """
+    for code_block_utf8, node in code_blocks_iter:
+        base_offset = node.start_byte
+        subblock_ranges = _extract_subblock_ranges(config, node, base_offset)
+        
+        if not subblock_ranges:
+            continue 
+
+        max_bytes_per_subblock = int(config.max_token_sequence_length * bytes_per_token_ratio)
+
+        subblock_ranges = _filter_subblocks(subblock_ranges, max_bytes_per_subblock) 
+        
+        code_block_utf8 = code_block_utf8[:max_bytes_per_subblock]  # Trunctate code block code if it is larger than bytes_per_code_block, we only consider subblocks inside this range.
+
+        fim_examples = _generate_fim_examples_from_code_block(config, code_block_utf8, subblock_ranges, bytes_per_token_ratio)
+        for fim_example in fim_examples:
+            yield fim_example
+
+
+def _save_tokenized_batch_as_jsonl(file_path: Path, batch: Mapping[str, List[List[int]]]) -> None:
+    with open(file_path, 'a', encoding='utf-8') as f:
+        batch_size = len(batch['input_ids'])
+        for i in range(batch_size):
+            input_ids = batch['input_ids'][i]
+            attention_mask = batch['attention_mask'][i]
+            labels = batch['labels'][i]
+            
+            example = {
+                'input_ids': input_ids,
+                'attention_mask': attention_mask,
+                'labels': labels
+            }
+            f.write(json.dumps(example, ensure_ascii=False) + '\n')  # ensre utf8 encoding
+
+
+def tokenize_and_save_fim_examples(config: Config, file_path: Path, fim_examples_iter: Iterator[bytes], tokenizer: AutoTokenizer) -> None: 
+    examples_counter = 0
+    batch = []
+    for fim_example in fim_examples_iter:
+        batch.append(fim_example.decode('utf-8'))
+        examples_counter += 1
+        if (len(batch) == config.tokenizer_batch_size):
+            tokenized_batch = tokenizer(
+                batch,
+                padding=False,
+                return_tensors=None,
+                return_attention_mask=True
+            )
+            tokenized_batch["labels"] = tokenized_batch["input_ids"]
+            _save_tokenized_batch_as_jsonl(file_path, tokenized_batch)
+            batch = []
+
+    # last batch is smaller than config.tokenizer_batch_size, the "rest"
+    if batch:
+        tokenized_batch = tokenizer(
+            batch,
+            padding=False,
+            return_tensors=None,
+            return_attention_mask=True
+        )
+        tokenized_batch["labels"] = tokenized_batch["input_ids"]
+        _save_tokenized_batch_as_jsonl(file_path, tokenized_batch)
+
+    logger.info(f"Processed and saved {examples_counter} FIM examples to {file_path}")

codefinetuner/preprocess/run.py

@@ -0,0 +1,133 @@
+import argparse
+import logging.config
+import sys
+from pathlib import Path
+
+from transformers import AutoTokenizer
+
+from .config import Config
+from .extract import get_code_blocks_from_auto_split, get_code_blocks_from_manual_split
+from .process import create_fim_examples, estimate_bytes_per_token_ratio, tokenize_and_save_fim_examples
+
+
+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="Preprocess code dataset for FIM fine-tuning.")
+    parser.add_argument(
+        "--config",
+        type=Path,
+        required=True,
+        help="YAML config file path with 'preprocess:' 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.",
+    )
+    return parser.parse_args()
+
+
+def _clear_existing_datasets(config: Config) -> None:
+    for file in [config.train_dataset_path, config.eval_dataset_path, config.test_dataset_path]:
+        if file.exists():
+            file.unlink()
+            logger.info(f"Deleted old dataset: {file}")
+
+
+def _validate_and_configure_tokenizer(config: Config, tokenizer: AutoTokenizer):
+    required_tokens = [
+        config.fim_prefix_token,
+        config.fim_middle_token,
+        config.fim_suffix_token,
+        config.fim_pad_token,
+        config.eos_token
+    ]
+
+    for token in required_tokens:
+        token_id = tokenizer.convert_tokens_to_ids(token)
+        
+        if token_id == tokenizer.unk_token_id or token_id is None:
+            raise ValueError(f"Token '{token}' defined in Config is missing from Tokenizer vocabulary.")
+
+    tokenizer.pad_token = config.fim_pad_token
+    tokenizer.eos_token = config.eos_token
+
+    logger.info("Tokenizer validated and configured successfully.")
+
+
+def run(config: Config) -> None:
+    _clear_existing_datasets(config)
+
+    if config.split_mode == "auto":
+        logger.info("Using auto-generated dataset split.")
+        train_code_blocks_iter, eval_code_blocks_iter, test_code_blocks_iter = get_code_blocks_from_auto_split(config) 
+    else: # config.split_mode == "manual":
+        logger.info("Using manual dataset split from directories.")
+        train_code_blocks_iter, eval_code_blocks_iter, test_code_blocks_iter = get_code_blocks_from_manual_split(config) 
+
+    tokenizer = AutoTokenizer.from_pretrained(config.model_name)
+    _validate_and_configure_tokenizer(config, tokenizer)
+
+    bytes_per_token_ratio = estimate_bytes_per_token_ratio(config, tokenizer, number_of_code_blocks=20000)
+    logger.info(f"Estimated bytes_per_token_ratio: {bytes_per_token_ratio}")
+
+    train_fim_examples_iter= create_fim_examples(config, train_code_blocks_iter, bytes_per_token_ratio)
+    eval_fim_examples_iter = create_fim_examples(config, eval_code_blocks_iter, bytes_per_token_ratio)
+    test_fim_examples_iter = create_fim_examples(config, test_code_blocks_iter, bytes_per_token_ratio)
+    
+    tokenize_and_save_fim_examples(config, config.train_dataset_path, train_fim_examples_iter, tokenizer)
+    tokenize_and_save_fim_examples(config, config.eval_dataset_path, eval_fim_examples_iter, tokenizer)
+    tokenize_and_save_fim_examples(config, config.test_dataset_path, test_fim_examples_iter, tokenizer)
+
+    logger.info("Saved train, eval, test datasets to disk")  
+
+
+def main() -> None:
+    user_args = _parse_args()
+    _setup_logger(user_args.log_level)
+    try:
+        preprocess_config = Config.load_from_yaml(user_args.config)
+        run(preprocess_config)
+    except Exception:
+        logger.exception(f"Preprocessing failed")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()
+    

codefinetuner/preprocess/tree_sitter_definitions.json

@@ -0,0 +1,189 @@
+{
+  "c": {
+    "block_types": [
+      "function_definition",
+      "struct_specifier",
+      "union_specifier",
+      "enum_specifier"
+    ],
+    "subblock_types": [
+      "compound_statement",
+      "parameter_list",
+      "declaration",
+      "expression_statement",
+      "if_statement",
+      "while_statement",
+      "for_statement",
+      "switch_statement",
+      "case_statement",
+      "return_statement",
+      "field_declaration_list",
+      "field_declaration",
+      "enumerator_list",
+      "enumerator"
+    ]
+  },
+  "python": {
+    "block_types": [
+      "function_definition",
+      "class_definition",
+      "decorated_definition",
+      "async_function"
+    ],
+    "subblock_types": [
+      "block",
+      "parameters",
+      "assignment",
+      "expression_statement",
+      "if_statement",
+      "while_statement",
+      "for_statement",
+      "match_statement",
+      "case_clause",
+      "return_statement",
+      "decorated_definition",
+      "try_statement",
+      "except_clause",
+      "finally_clause",
+      "with_statement",
+      "lambda",
+      "yield"
+    ]
+  },
+  "cpp": {
+    "block_types": [
+      "function_definition",
+      "class_specifier",
+      "struct_specifier",
+      "union_specifier",
+      "enum_specifier",
+      "template_declaration",
+      "namespace_definition"
+    ],
+    "subblock_types": [
+      "compound_statement",
+      "parameter_list",
+      "declaration",
+      "expression_statement",
+      "if_statement",
+      "while_statement",
+      "for_statement",
+      "switch_statement",
+      "case_statement",
+      "return_statement",
+      "field_declaration_list",
+      "field_declaration",
+      "enumerator_list",
+      "enumerator",
+      "try_statement",
+      "catch_statement",
+      "finally_clause"
+    ]
+  },
+  "java": {
+    "block_types": [
+      "method_declaration",
+      "class_declaration",
+      "interface_declaration",
+      "enum_declaration",
+      "record_declaration",
+      "constructor_declaration"
+    ],
+    "subblock_types": [
+      "block",
+      "formal_parameters",
+      "local_variable_declaration",
+      "expression_statement",
+      "if_statement",
+      "while_statement",
+      "for_statement",
+      "enhanced_for_statement",
+      "switch_expression",
+      "switch_rule",
+      "case_label",
+      "return_statement",
+      "throw_statement",
+      "catch_clause",
+      "finally_clause",
+      "yield_statement",
+      "try_statement",
+      "try_with_resources_statement",
+      "synchronized_statement",
+      "static_initializer"
+    ]
+  }, 
+  "vhdl": {
+    "block_types": [
+      "entity_declaration",
+      "architecture_body",
+      "process_statement",
+      "package_declaration",
+      "configuration_declaration",
+      "generate_statement"
+    ],
+    "subblock_types": [
+      "signal_assignment",
+      "variable_assignment",
+      "if_statement",
+      "case_statement",
+      "loop_statement",
+      "wait_statement",
+      "assert_statement",
+      "component_instantiation",
+      "concurrent_statement",
+      "sequential_statement"
+    ]
+  },
+  "bash": {
+    "block_types": [
+      "function_definition",
+      "compound_statement",
+      "subshell",
+      "pipeline",
+      "list"
+    ],
+    "subblock_types": [
+      "command",
+      "variable_assignment",
+      "redirected_statement",
+      "expression",
+      "assignment",
+      "test_command",
+      "negated_command",
+      "break_statement",
+      "continue_statement",
+      "return_statement",
+      "command_substitution",
+      "process_substitution",
+      "arithmetic_expansion",
+      "if_statement",
+      "while_statement",
+      "for_statement",
+      "c_style_for_statement",
+      "case_statement"
+    ]
+  },
+  "mojo": {
+    "block_types": [
+      "function_definition",
+      "struct_definition",
+      "trait_definition",
+      "class_definition",
+      "enum_definition"
+    ],
+    "subblock_types": [
+      "compound_statement",
+      "parameter_list",
+      "declaration",
+      "expression_statement",
+      "return_statement",
+      "assignment",
+      "augmented_assignment",
+      "try_statement",
+      "catch_clause",
+      "finally_clause",
+      "decorated_definition"
+    ]
+  }
+}
+