codefinetuner 0.1.0__py3-none-any.whl

codefinetuner/evaluate/generate.py

@@ -0,0 +1,145 @@
+import gc
+import json
+import logging
+import math
+from pathlib import Path
+
+import torch
+from peft import PeftModel
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from .config import Config
+
+
+logger = logging.getLogger(__name__)
+
+
+def _load_lora_model(config: Config, checkpoint_path: Path) ->AutoModelForCausalLM:
+    # load base model to CPU first to prevent VRAM fragmentation/OOM
+    base_model = AutoModelForCausalLM.from_pretrained(
+        pretrained_model_name_or_path=config.model_name,
+        dtype=config.model_dtype,
+        low_cpu_mem_usage=True
+    )
+
+    # load and attach LoRA adapter and move to device (CUDA, MPS or remains on CPU)
+    lora_model = PeftModel.from_pretrained(
+        model=base_model, 
+        model_id=checkpoint_path
+    ).to(config.device)
+
+    lora_model.eval()
+
+    logger.info(f"Loaded LoRA model from checkpoint {checkpoint_path} to device {config.device}")
+    return lora_model
+
+
+def _load_tokenizer(config: Config) -> AutoTokenizer:
+    tokenizer = AutoTokenizer.from_pretrained(config.model_name)
+    tokenizer.pad_token = config.fim_pad_token
+    tokenizer.padding_side = "right"
+    return tokenizer
+
+
+def _get_fim_perplexity(config: Config, model: AutoModelForCausalLM, 
+                        perplexity_input_token_ids: list[int], perplexity_label_token_ids: list[int]) -> float:
+    """
+    FIM perplexity: Measures model confidence in the ground truth reference middle code.
+    (How surprised is the model by the ground truth reference middle code).
+    Lower perplexity indicates higher confidence. perplexity = exp(loss).
+    """
+    try:
+        input_tensor = torch.tensor([perplexity_input_token_ids], device=config.device)
+        label_tensor = torch.tensor([perplexity_label_token_ids], device=config.device)
+        with torch.inference_mode():
+            outputs = model(input_ids=input_tensor, 
+                            labels=label_tensor)
+            loss = outputs.loss
+            
+        return math.exp(loss.item())
+        
+    except Exception as e:
+        logger.warning(f"Perplexity calcualtin failed, returning inf: {e}")
+        return float('inf')
+
+
+def _generate(config: Config, model: AutoModelForCausalLM, tokenizer: AutoTokenizer, prompt_token_ids: list[int]) -> list:
+    model.eval()
+    prompt_token_ids_tensor = torch.tensor([prompt_token_ids], device=config.device)
+    with torch.inference_mode():
+        generated_token_ids_tensor = model.generate(
+            input_ids=prompt_token_ids_tensor,
+            max_new_tokens=config.generation_max_new_tokens,
+            do_sample=config.generation_do_sample,
+            temperature=config.generation_temperature,
+            top_p=config.generation_top_p,
+            pad_token_id=tokenizer.pad_token_id
+        )
+
+        # slice the output, take everything after the input_length, model.generate() functin returns the whole example, not only the generated text
+        generated_middle_token_ids_tensor = generated_token_ids_tensor[0][prompt_token_ids_tensor.shape[1] :]
+        generated_middle_token_ids = generated_middle_token_ids_tensor.tolist()
+
+    return generated_middle_token_ids
+
+
+def _clear_hardware_cache(config: Config) -> None:
+    gc.collect()
+    if config.device == "cuda":
+        torch.cuda.synchronize()
+        torch.cuda.empty_cache()
+    elif config.device == "mps":
+        torch.mps.empty_cache()
+
+
+def generate_and_save(config: Config, checkpoint_path: Path):
+    lora_model = _load_lora_model(config, checkpoint_path)
+    tokenizer = _load_tokenizer(config)
+    fim_prefix_token_id = tokenizer.convert_tokens_to_ids(config.fim_prefix_token)
+    fim_suffix_token_id = tokenizer.convert_tokens_to_ids(config.fim_suffix_token)
+    fim_middle_token_id = tokenizer.convert_tokens_to_ids(config.fim_middle_token)
+    
+    line_counter = 0
+    try:
+        with config.benchmark_dataset_path.open("r") as benchmark_dataset_file, \
+            config.benchmark_evaluation_results_path.open("w") as evaluation_results_file:
+            
+            for line in benchmark_dataset_file:
+                benchmark_example = json.loads(line)
+                prompt_token_ids = ([fim_prefix_token_id] + benchmark_example["prefix_token_ids"] +
+                                   [fim_suffix_token_id] + benchmark_example["suffix_token_ids"] +
+                                   [fim_middle_token_id])
+                
+                perplexity_input_token_ids = benchmark_example["example_token_ids"]
+                perplexity_label_token_ids = perplexity_input_token_ids.copy()
+                perplexity_label_token_ids[:len(prompt_token_ids)] = [config.label_pad_token_id] * len(prompt_token_ids)  # mask labels all except ground truth reference middle tokens
+
+                lora_generated_middle_token_ids = _generate(config, lora_model, tokenizer, prompt_token_ids)
+                lora_perplexity = _get_fim_perplexity(config, lora_model, perplexity_input_token_ids, perplexity_label_token_ids)
+
+                with lora_model.disable_adapter():
+                    base_generated_middle_token_ids = _generate(config, lora_model, tokenizer, prompt_token_ids)
+                    base_perplexity = _get_fim_perplexity(config, lora_model, perplexity_input_token_ids, perplexity_label_token_ids)
+
+                reference_middle = benchmark_example["middle"]
+                lora_generated_middle = tokenizer.decode(lora_generated_middle_token_ids, skip_special_tokens=True)
+                base_generated_middle = tokenizer.decode(base_generated_middle_token_ids, skip_special_tokens=True)
+                result = {
+                    "example_id": line_counter,
+                    "reference_middle": reference_middle,
+                    "base_generated_middle": base_generated_middle,
+                    "lora_generated_middle": lora_generated_middle,
+                    "base_perplexity": base_perplexity,
+                    "lora_perplexity": lora_perplexity
+                }
+                evaluation_results_file.write(json.dumps(result) + "\n")
+
+                line_counter += 1
+                if line_counter % 10 == 0:
+                    _clear_hardware_cache(config)
+                    logger.info(f"Processed {line_counter} benchmark examples")
+        
+    except Exception as e:
+        raise RuntimeError(f"Generation failed at example {line_counter}: {e}") from e
+    
+    logger.info(f"Successfully generated and saved {line_counter} number of examples to {config.benchmark_evaluation_results_path}.")

codefinetuner/evaluate/metrics.py

@@ -0,0 +1,160 @@
+import logging
+import re
+from nltk.tokenize import word_tokenize
+from nltk.translate.bleu_score import SmoothingFunction, sentence_bleu
+
+from .codebleu_shim import codebleu_score
+from .config import Config
+
+
+logger = logging.getLogger(__name__)
+
+
+def _codebleu_structure_valid(config: Config, reference: str) -> bool:
+    """
+    Checks if the reference code is structurally complex enough for CodeBLEU.
+    Performs a self-match test focusing only on Syntax (AST) and Data-flow. 
+    If the reference is too simple (e.g., import blocks, comment) 
+    to build these structures, the example is skipped. 
+    """
+    # suppress root logger warnings that are raised if it is not possible to calculate all 4 metrics
+    root_logger = logging.getLogger()
+    original_level = root_logger.getEffectiveLevel() 
+    root_logger.setLevel(logging.ERROR)
+    
+    try:
+        test_weights = (0.0, 0.0, 0.5, 0.5) 
+        result = codebleu_score([reference], [reference], 
+                               lang=config.codebleu_language, 
+                               weights=test_weights)
+    
+        syntax_valid = result.get('syntax_match_score', 0) > 0
+        dataflow_valid = result.get('dataflow_match_score', 0) > 0
+        return syntax_valid and dataflow_valid
+    except Exception:
+        return False
+    finally:
+        # restore logger
+        root_logger.setLevel(original_level)
+
+
+def get_codebleu(config: Config, reference: str, prediction: str) -> tuple[float, bool]:
+    """
+    CodeBLEU: Computes weighted combination of four different similarity metrics:
+    1. N-gram Match: Standard surface-level text overlap.
+    2. Weighted N-gram: Text overlap with priority on keywords (if, else, etc.).
+    3. Syntax (AST): Structural similarity between Abstract Syntax Trees.
+    4. Data-flow: Logical similarity based on variable dependencies and usage.
+
+    Examples are skipped if structural components (AST/Data-flow) cannot be 
+    extracted from the reference. This prevents the final average score from 
+    being unfairly lowered by snippets that cannot be properly parsed, 
+    ensuring a more accurate representation of model quality.
+
+    Standard weights are: [0.25, 0.25, 0.25, 0.25].
+    CodeBLEU = (codebleu_ngram_weight * ngram_score) + 
+               (codebleu_weighted_ngram_weight * weighted_ngram_score) + 
+               (codebleu_syntax_ast_weight * syntax_ast_score) + 
+               (codebleu_dataflow_weight * dataflow_score)
+    """
+    if not _codebleu_structure_valid(config, reference):
+        return (0.0, False)
+        
+    try:
+        codebleu_algorithm_weights = (
+            config.codebleu_ngram_weight, 
+            config.codebleu_weighted_ngram_weight, 
+            config.codebleu_syntax_ast_weight, 
+            config.codebleu_dataflow_weight
+        )
+        
+        result = codebleu_score(
+            [reference], [prediction], 
+            lang=config.codebleu_language, 
+            weights=codebleu_algorithm_weights
+        )
+        
+        return (float(result['codebleu']), True)
+        
+    except Exception as e:
+        logger.warning(f"CodeBLEU calculation failed, returning (0.0, False): {e}")
+        return (0.0, False)
+
+
+def get_sentencebleu(config: Config, reference: str, prediction: str) -> float:
+    """
+    SentenceBLEU: Measures n-gram overlap between reference and prediction.
+    It rewards matching sequences of words (1-4) and uses smoothing 
+    (Method1: Adds a tiny epsilon to all n-gram counts) to prevent a total 
+    0.0 score when long sequences (e.g. 4-grams) don't match exactly.
+    """
+    config.ensure_nltk_initialized()  # make sure that required nltk downloads are done
+
+    try:
+        reference_tokens = word_tokenize(reference)
+        prediction_tokens = word_tokenize(prediction)
+        
+        weights = (
+            config.sentencebleu_ngram_weight_1, 
+            config.sentencebleu_ngram_weight_2, 
+            config.sentencebleu_ngram_weight_3, 
+            config.sentencebleu_ngram_weight_4
+        )
+        
+        smoothing = SmoothingFunction().method1
+        
+        score = sentence_bleu(
+            [reference_tokens], 
+            prediction_tokens, 
+            weights=weights, 
+            smoothing_function=smoothing
+        )
+        
+        return float(score)
+        
+    except Exception as e:
+        logger.warning(f"SentenceBLEU calculation failed, returning 0.0: {e}")
+        return 0.0
+
+
+def get_exact_match(reference: str, prediction: str) -> float:
+    """Exact match is 1.0 if identical, 0.0 otherwise. Collapese all whitespaces."""
+    try:
+        # re.sub(r'\s+', ' ', text.strip()): Collapses whitespace to compare logic regardless of formatting.
+        ref_norm = re.sub(r'\s+', ' ', reference.strip())
+        pred_norm = re.sub(r'\s+', ' ', prediction.strip())
+        
+        if ref_norm == pred_norm:
+            return 1.0
+        else:
+            return 0.0
+    except Exception as e:
+        logger.warning(f"Exact match calculation failed, returning 0.0: {e}")
+        return 0.0
+
+
+def get_line_match(config: Config, reference: str, prediction: str) -> float:
+    """Check if the first n lines match, ignoring trailing whitespace."""
+    try:
+        n = config.line_match_number_of_lines
+
+        # line.rstrip(): Removes trailing whitespace while preserving leading indentation.
+        ref_lines_stripped = []
+        for line in reference.splitlines()[:n]:
+            ref_lines_stripped.append(line.rstrip())
+
+        pred_lines_stripped = []
+        for line in prediction.splitlines()[:n]:
+            pred_lines_stripped.append(line.rstrip())
+
+        # Ensure both lists have the required number of lines
+        if len(pred_lines_stripped) < n or len(ref_lines_stripped) < n:
+            return 0.0
+
+        if pred_lines_stripped == ref_lines_stripped:
+            return 1.0
+        else:
+            return 0.0
+    except Exception as e:
+        logger.warning(f"Line match calculation failed, returning 0.0: {e}")
+        return 0.0

codefinetuner/evaluate/run.py

@@ -0,0 +1,125 @@
+import argparse
+import logging
+import sys
+import logging.config
+from pathlib import Path
+
+from transformers.trainer_utils import get_last_checkpoint
+
+from .config import Config
+from .benchmark import create_benchmark_dataset
+from .generate import generate_and_save
+from .evaluate import evaluate_and_save
+from .analyze import analyze_metric, save_all_metric_stats, get_plot_path, plot_metric_and_save, plot_all_metric_averages_and_save
+
+
+logger = logging.getLogger(__name__)
+
+
+def _parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Run evaluation")
+    parser.add_argument(
+        "--config",
+        type=Path,
+        required=True,
+        help="Path to the pipeline YAML configuration file.",
+    )
+    parser.add_argument(
+        "--log-level",
+        type=str,
+        default="INFO",
+        choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
+        help="Logging level.",
+    )
+    return parser.parse_args()
+
+
+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 _silence_noisy_third_party_loggers() -> None:
+    """Reduce log level of known noisy third‑party loggers."""
+    logging.getLogger("matplotlib").setLevel(logging.WARNING)
+    logging.getLogger("matplotlib.font_manager").setLevel(logging.WARNING)
+    logging.getLogger("nltk").setLevel(logging.WARNING)
+    # add more as needed
+
+
+def _ensure_checkpoints(config: Config) -> None:
+    checkpoints_dir = config.trainer_checkpoints_dir_path
+    if not checkpoints_dir.exists():
+        raise RuntimeError("Checkpoint directory not found.")
+    if not list(checkpoints_dir.iterdir()):
+        raise RuntimeError("No checkpoints in directory.")
+
+
+def run(config: Config) -> None:
+    if not config.benchmark_use_existing_dataset or not config.benchmark_dataset_path.exists():
+        dataset_len = create_benchmark_dataset(config)
+        logger.info(f"Created new benchmark dataset '{config.benchmark_dataset_path}' with '{dataset_len}' examples")
+    else:
+        logger.info(f"Proceeding with existing file '{config.benchmark_dataset_path}'...")
+    
+    if not config.plot_only: 
+        _ensure_checkpoints(config)
+        if config.trainer_checkpoint == "last":
+            checkpoint_path = get_last_checkpoint(config.trainer_checkpoints_dir_path)
+        else:
+            checkpoint_path = config.trainer_checkpoints_dir_path / config.trainer_checkpoint 
+        generate_and_save(config, checkpoint_path)
+        evaluate_and_save(config)
+
+    all_metric_stats_np = []
+    for metric_name, higher_is_better in config.metric_configs:
+        metric_stats_np = analyze_metric(config, metric_name, higher_is_better)
+        plot_path = get_plot_path(config.benchmark_evaluation_results_dir, metric_name)
+        plot_metric_and_save(metric_stats_np, metric_name, plot_path)
+        all_metric_stats_np.append(metric_stats_np)
+    
+    all_metric_averages_plot_path = get_plot_path(config.benchmark_evaluation_results_dir, "all_metric_averages")
+    plot_all_metric_averages_and_save(all_metric_stats_np, all_metric_averages_plot_path) 
+    save_all_metric_stats(config, all_metric_stats_np)
+
+
+def main():
+    user_args = _parse_args()
+    _setup_logger(user_args.log_level)
+    _silence_noisy_third_party_loggers()
+    try:
+        evaluate_config = Config.load_from_yaml(user_args.config)
+        run(evaluate_config)
+    except Exception:
+        logger.exception("Evaluation failed")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

codefinetuner/finetune/config.py

@@ -0,0 +1,186 @@
+import logging
+import math
+from dataclasses import dataclass, field, fields
+from pathlib import Path
+from typing import List, Any
+
+import torch
+from omegaconf import OmegaConf, MISSING
+
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class Config:
+# --- Mandatory Parameters ---
+    model_name: str = MISSING
+    fim_pad_token: str = MISSING
+    label_pad_token_id: int = MISSING
+
+    # --- Model Settings ---
+    model_attn_implementation: str = "sdpa"  # sdpa = built-in PyTorch implementation of scaled dot product attention, imporves performance and memory efficiency
+    model_dtype: Any = field(init=False)    # changed to Any to bypass OmegaConf validation
+
+    # --- LoRA Settings ---
+    lora_r: int = 32
+    lora_alpha: int = 64
+    lora_dropout: float = 0.1
+    lora_bias: str = "none"
+    lora_target_modules: List[str] = field(default_factory=lambda: [
+        "q_proj", "v_proj", "k_proj", "o_proj",
+        "gate_proj", "down_proj", "up_proj"
+    ])
+
+    # --- Trainer Hyperparameters ---
+    trainer_resume_from_checkpoint: str | None = "last"
+    trainer_clear_checkpoint_dir: bool = False
+    trainer_num_train_epochs: int = 1
+    trainer_per_device_train_batch_size: int = 2
+    trainer_per_device_eval_batch_size: int = 2
+    trainer_gradient_accumulation_steps: int = 32  # simulates a larger effective batch size, number of forward/backward passes to accumulate before performing one optimizer step
+    trainer_learning_rate: float = 2e-5
+    trainer_weight_decay: float = 0.1
+    trainer_max_grad_norm: float = 1.0
+    trainer_lr_scheduler_type: str = "cosine"
+    trainer_warmup_steps: int = 50
+    trainer_gradient_checkpointing: bool = True  # saves memory by storing only key "checkpoint" activations and re-calculating intermediate ones during the backward pass
+    trainer_max_steps: int = field(init=False)
+
+    # --- Logging and Evaluation Strategy ---
+    trainer_logging_steps: int = 10  # average training loss over trainer_logging_steps period is calculated and logged
+    trainer_eval_strategy: str = "steps"
+    trainer_eval_steps: int = 100
+    trainer_save_strategy: str = "steps"
+    trainer_save_steps: int = 100
+    trainer_logging_strategy: str = "steps"
+
+    # --- Dataset ---
+    dataset_shuffle_buffer_size: int = 50000
+    dataset_shuffle_seed: int = 0
+    dataset_train_dataset_length: int = field(init=False)
+
+    # --- Hardware Configuration --- 
+    device: str = field(init=False)
+    
+    # --- Path Management ---
+    workspace_path: Path | None = None 
+    train_dataset_path: Path = field(init=False)
+    eval_dataset_path: Path = field(init=False)
+    finetune_outputs_dir_path: Path = field(init=False)
+    trainer_checkpoints_dir_path: Path = field(init=False)
+    trainer_log_path: Path = field(init=False)
+    trainer_plot_path: Path = field(init=False)
+    lora_adapter_path: Path = field(init=False)
+    lora_model_path: Path = field(init=False)
+    trainer_model_merge_offload_folder_path: Path = field(init=False)
+
+    @classmethod
+    def load_from_yaml(cls, yaml_path: Path) -> "Config":
+        if not yaml_path.exists():
+            raise FileNotFoundError(f"Config file not found: {yaml_path}")
+        
+        logger.info(f"Loading configuration from {yaml_path}") 
+        config_dict = OmegaConf.structured(cls)
+        try:
+            yaml_file_node = OmegaConf.load(yaml_path)
+        except Exception as e:
+            raise ValueError(f"Failed to load YAML config {yaml_path}") from e
+
+        yaml_file_dict = OmegaConf.to_container(yaml_file_node, resolve=True)
+        yaml_finetune_dict = yaml_file_dict.get("finetune", {})
+
+        yaml_finetune_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_finetune_dict:
+                yaml_finetune_valid_dict[field.name] = yaml_finetune_dict[field.name]
+        logger.debug(f"Filtered YAML configuration: {yaml_finetune_valid_dict}")
+
+        merged_config_dict = OmegaConf.merge(config_dict, yaml_finetune_valid_dict)
+        return OmegaConf.to_object(merged_config_dict)
+
+    def __post_init__(self) -> None:
+        self._setup_device_and_precision()
+        self._setup_paths()
+        self._ensure_output_paths_exist()
+        self.dataset_train_dataset_length = self._get_dataset_length(self.train_dataset_path)
+        self.trainer_max_steps = self._calculate_max_steps()
+        
+    def _setup_device_and_precision(self) -> None:
+        if torch.cuda.is_available():
+            self.device = "cuda"
+            self.model_dtype = torch.bfloat16
+        elif torch.backends.mps.is_available():
+            self.device = "mps"
+            self.model_dtype = torch.float16
+        else:
+            self.device = "cpu"
+            self.model_dtype = torch.float32
+        logger.info(f"Execution environment: device={self.device}, dtype={self.model_dtype}")
+    
+    def _setup_paths(self) -> None:
+        if self.workspace_path is None:
+            self.workspace_path = Path.cwd()
+        self.train_dataset_path = self.workspace_path / "outputs" / "preprocess" / "results" / "datasets" / "train_dataset.jsonl"
+        self.eval_dataset_path = self.workspace_path / "outputs" / "preprocess" / "results" / "datasets" / "eval_dataset.jsonl"
+        self.finetune_outputs_dir_path = self.workspace_path / "outputs" / "finetune"
+        self.trainer_checkpoints_dir_path = self.finetune_outputs_dir_path / "checkpoints"
+        self.trainer_model_merge_offload_folder_path = self.finetune_outputs_dir_path / "trainer_model_merge_offload_folder"
+        self.trainer_log_path = self.finetune_outputs_dir_path / "results" / "trainer_log.json" 
+        self.trainer_plot_path = self.finetune_outputs_dir_path / "results" / "trainer_loss_plot.png"
+        self.lora_adapter_path = self.finetune_outputs_dir_path / "results" / "lora_adapter"
+        self.lora_model_path = self.finetune_outputs_dir_path / "results" / "lora_model"
+        logger.debug(f"Resolved workspace path to: {self.workspace_path}")
+    
+    def _ensure_output_paths_exist(self) -> None:
+        paths = [
+            self.finetune_outputs_dir_path,
+            self.trainer_checkpoints_dir_path,
+            self.trainer_model_merge_offload_folder_path,
+            self.trainer_log_path, 
+            self.trainer_plot_path,
+            self.lora_adapter_path,
+            self.lora_model_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 _get_dataset_length(self, path: Path) -> int:
+        """
+        Calculates line count for streaming dataset progress estimation.
+        We need to do this because we load it as a streaming dataset iterator, which can only be iterated once.
+        """
+        if not path.exists():
+            raise FileNotFoundError( f"Training dataset not found at expected path: {path}. Ensure the dataset file exists before initializing the Config.")
+        count = 0
+        with open(path, "r", encoding="utf-8") as file:
+            for _ in file:
+                count += 1
+        logger.debug(f"Dataset length for '{path.name}': {count} lines")
+        return count
+        
+    def _calculate_max_steps(self) -> int:
+        """
+        Calculates total steps based on effective batch size and dataset length.
+        Because we use streaming dataset iterators for efficiency, we cannot use num_train_epochs trainer class parameter directly.
+        Instead, we need to calculate max_steps and pass it to the trainer.
+        """
+        if self.dataset_train_dataset_length == 0:
+            logger.warning("Dataset length is 0; max_steps will be 0.")
+            return 0
+            
+        effective_batch_size = (self.trainer_per_device_train_batch_size * self.trainer_gradient_accumulation_steps)
+        if effective_batch_size == 0:
+            raise ValueError("Effective batch size (batch_size * grad_accum) cannot be zero.")
+        steps_per_epoch = math.ceil(self.dataset_train_dataset_length / effective_batch_size)
+        max_steps = steps_per_epoch * self.trainer_num_train_epochs
+        logger.debug(f"Calculated training schedule: {max_steps} total steps ({steps_per_epoch} steps/epoch for {self.trainer_num_train_epochs} epochs)")
+        return max_steps
+       

codefinetuner/finetune/model.py

@@ -0,0 +1,59 @@
+
+import logging
+
+from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
+from transformers import AutoModelForCausalLM, BitsAndBytesConfig
+
+from .config import Config
+
+
+logger = logging.getLogger(__name__)
+
+
+def load_and_configure_lora_model(config: Config) -> AutoModelForCausalLM:
+    if config.device == "cuda":
+        bnb_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",  # nf4 (NormalFloat4) = optimal for LLM training, because llm weights follwo normal distribution  
+            bnb_4bit_compute_type=config.model_dtype,
+            bnb_4bit_use_double_quant=True,
+        )
+        model = AutoModelForCausalLM.from_pretrained(
+            pretrained_model_name_or_path=config.model_name,
+            attn_implementation=config.model_attn_implementation,  
+            quantization_config=bnb_config,
+            device_map="auto" # let bitsandbytes handle placement
+        )
+        model = prepare_model_for_kbit_training(model)
+    elif config.device == "mps":
+        model = AutoModelForCausalLM.from_pretrained(
+            pretrained_model_name_or_path=config.model_name,
+            attn_implementation=config.model_attn_implementation,
+            dtype=config.model_dtype 
+        ).to("mps")
+    else:
+        model = AutoModelForCausalLM.from_pretrained(
+            pretrained_model_name_or_path=config.model_name,
+            attn_implementation=config.model_attn_implementation,
+            dtype=config.model_dtype
+        ).to("cpu")
+    
+    logger.info(
+        f"Model: {config.model_name} | Device: {config.device} | "
+        f"Dtype: {config.model_dtype} | Attn: {config.model_attn_implementation}"
+    )
+    
+    # forces the input to require gradients, ensuring the backward pass graph stays connected when using frozen base models with gradient checkpointing
+    if config.trainer_gradient_checkpointing:
+        model.enable_input_require_grads()
+
+    lora_config = LoraConfig(
+        lora_alpha=config.lora_alpha, 
+        lora_dropout=config.lora_dropout, 
+        r=config.lora_r, 
+        bias=config.lora_bias, 
+        task_type="CAUSAL_LM",
+        target_modules=config.lora_target_modules,
+    )
+    lora_model = get_peft_model(model, lora_config)
+    return lora_model