Retrospective: 2 Years of Using Hugging Face Transformers 4.40 for Enterprise Code Fine-Tuning

import os
import sys
import json
import logging
from dataclasses import dataclass, field
from typing import Optional, Dict, List

import torch
from torch.utils.data import Dataset, DataLoader
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    Trainer,
    TrainingArguments,
    HfArgumentParser,
    set_seed,
    EarlyStoppingCallback
)
from transformers.utils import is_flash_attention_available
import datasets

# Configure logging for audit trails
logging.basicConfig(
    format=\"%(asctime)s - %(levelname)s - %(name)s - %(message)s\",
    level=logging.INFO,
    handlers=[logging.StreamHandler(sys.stdout)]
)
logger = logging.getLogger(__name__)

@dataclass
class FineTuneArgs:
    model_path: str = field(default=\"codellama/CodeLlama-7b-Instruct-hf\", metadata={\"help\": \"Base model to fine-tune\"})
    train_data_path: str = field(default=\"enterprise_code_train.jsonl\", metadata={\"help\": \"Path to training data JSONL\"})
    eval_data_path: str = field(default=\"enterprise_code_eval.jsonl\", metadata={\"help\": \"Path to eval data JSONL\"})
    output_dir: str = field(default=\"./codellama-7b-enterprise-finetuned\", metadata={\"help\": \"Output directory for checkpoints\"})
    max_seq_length: int = field(default=2048, metadata={\"help\": \"Maximum sequence length for tokenization\"})
    use_flash_attention: bool = field(default=True, metadata={\"help\": \"Enable FlashAttention-2 if available\"})

class EnterpriseCodeDataset(Dataset):
    def __init__(self, data_path: str, tokenizer: AutoTokenizer, max_seq_length: int):
        self.tokenizer = tokenizer
        self.max_seq_length = max_seq_length
        self.examples = []

        # Load and validate training data
        if not os.path.exists(data_path):
            raise FileNotFoundError(f\"Training data not found at {data_path}\")
        try:
            with open(data_path, \"r\") as f:
                for line_num, line in enumerate(f, 1):
                    try:
                        example = json.loads(line.strip())
                        if \"input\" not in example or \"output\" not in example:
                            logger.warning(f\"Skipping line {line_num} in {data_path}: missing input/output keys\")
                            continue
                        self.examples.append(example)
                    except json.JSONDecodeError as e:
                        logger.error(f\"JSON parse error on line {line_num} in {data_path}: {e}\")
                        continue
        except Exception as e:
            logger.error(f\"Failed to load data from {data_path}: {e}\")
            raise

        logger.info(f\"Loaded {len(self.examples)} valid examples from {data_path}\")

    def __len__(self):
        return len(self.examples)

    def __getitem__(self, idx):
        example = self.examples[idx]
        # Format as instruction-response pair for CodeLlama Instruct
        prompt = f\"<>You are an enterprise Java developer. Write code for the following task.<>[INST]{example['input']}[/INST]\"
        full_text = prompt + example[\"output\"]

        # Tokenize with truncation and padding
        tokenized = self.tokenizer(
            full_text,
            max_length=self.max_seq_length,
            truncation=True,
            padding=\"max_length\",
            return_tensors=\"pt\"
        )

        # Create labels for causal LM (ignore prompt tokens in loss)
        labels = tokenized[\"input_ids\"].clone()
        prompt_tokenized = self.tokenizer(prompt, return_tensors=\"pt\")[\"input_ids\"]
        labels[:, :prompt_tokenized.shape[1]] = -100

        return {
            \"input_ids\": tokenized[\"input_ids\"].squeeze(0),
            \"attention_mask\": tokenized[\"attention_mask\"].squeeze(0),
            \"labels\": labels.squeeze(0)
        }

def main():
    parser = HfArgumentParser((FineTuneArgs, TrainingArguments))
    if len(sys.argv) == 2 and sys.argv[1].endswith(\".json\"):
        fine_tune_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
    else:
        fine_tune_args, training_args = parser.parse_args_into_dataclasses()

    set_seed(training_args.seed)

    # Validate FlashAttention availability for Transformers 4.40
    if fine_tune_args.use_flash_attention:
        if not is_flash_attention_available():
            logger.warning(\"FlashAttention-2 not available, falling back to standard attention\")
            fine_tune_args.use_flash_attention = False
        else:
            logger.info(\"FlashAttention-2 enabled for reduced VRAM usage\")

    # Load tokenizer with legacy=False for Transformers 4.40 compatibility
    try:
        tokenizer = AutoTokenizer.from_pretrained(
            fine_tune_args.model_path,
            legacy=False,
            padding_side=\"right\",
            truncation_side=\"right\"
        )
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
    except Exception as e:
        logger.error(f\"Failed to load tokenizer from {fine_tune_args.model_path}: {e}\")
        sys.exit(1)

    # Load model with FlashAttention if enabled
    try:
        model = AutoModelForCausalLM.from_pretrained(
            fine_tune_args.model_path,
            torch_dtype=torch.bfloat16,
            attn_implementation=\"flash_attention_2\" if fine_tune_args.use_flash_attention else \"eager\",
            device_map=\"auto\"
        )
    except Exception as e:
        logger.error(f\"Failed to load model from {fine_tune_args.model_path}: {e}\")
        sys.exit(1)

    # Load datasets
    try:
        train_dataset = EnterpriseCodeDataset(
            fine_tune_args.train_data_path,
            tokenizer,
            fine_tune_args.max_seq_length
        )
        eval_dataset = EnterpriseCodeDataset(
            fine_tune_args.eval_data_path,
            tokenizer,
            fine_tune_args.max_seq_length
        )
    except Exception as e:
        logger.error(f\"Failed to load datasets: {e}\")
        sys.exit(1)

    # Initialize Trainer with early stopping (Transformers 4.40 callback support)
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        callbacks=[EarlyStoppingCallback(early_stopping_patience=3)]
    )

    # Start fine-tuning
    logger.info(\"Starting fine-tuning...\")
    try:
        train_result = trainer.train()
        trainer.save_model(fine_tune_args.output_dir)
        trainer.log_metrics(\"train\", train_result.metrics)
        trainer.save_metrics(\"train\", train_result.metrics)
    except Exception as e:
        logger.error(f\"Fine-tuning failed: {e}\")
        sys.exit(1)

    # Evaluate
    logger.info(\"Running evaluation...\")
    try:
        eval_result = trainer.evaluate()
        trainer.log_metrics(\"eval\", eval_result)
        trainer.save_metrics(\"eval\", eval_result)
    except Exception as e:
        logger.error(f\"Evaluation failed: {e}\")
        sys.exit(1)

if __name__ == \"__main__\":
    main()

import os
import sys
import json
import logging
import argparse
from typing import List, Dict
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from datasets import load_dataset
import evaluate
import numpy as np

# Configure logging
logging.basicConfig(
    format=\"%(asctime)s - %(levelname)s - %(name)s - %(message)s\",
    level=logging.INFO,
    handlers=[logging.StreamHandler(sys.stdout)]
)
logger = logging.getLogger(__name__)

# Load evaluation metrics
code_eval = evaluate.load(\"code_eval\")
edit_sim = evaluate.load(\"edit_similarity\")

def parse_args():
    parser = argparse.ArgumentParser(description=\"Evaluate fine-tuned code model with Transformers 4.40\")
    parser.add_argument(\"--model_path\", type=str, required=True, help=\"Path to fine-tuned model\")
    parser.add_argument(\"--eval_data\", type=str, default=\"enterprise_code_eval.jsonl\", help=\"Evaluation data path\")
    parser.add_argument(\"--max_new_tokens\", type=int, default=512, help=\"Max new tokens to generate\")
    parser.add_argument(\"--batch_size\", type=int, default=4, help=\"Inference batch size\")
    parser.add_argument(\"--output_path\", type=str, default=\"eval_results.json\", help=\"Path to save results\")
    return parser.parse_args()

def load_eval_data(data_path: str) -> List[Dict]:
    \"\"\"Load and validate evaluation dataset\"\"\"
    examples = []
    if not os.path.exists(data_path):
        raise FileNotFoundError(f\"Eval data not found at {data_path}\")
    try:
        with open(data_path, \"r\") as f:
            for line_num, line in enumerate(f, 1):
                try:
                    example = json.loads(line.strip())
                    required_keys = [\"input\", \"expected_output\", \"test_cases\"]
                    if not all(k in example for k in required_keys):
                        logger.warning(f\"Skipping line {line_num}: missing required keys\")
                        continue
                    examples.append(example)
                except json.JSONDecodeError as e:
                    logger.error(f\"JSON error on line {line_num}: {e}\")
                    continue
    except Exception as e:
        logger.error(f\"Failed to load eval data: {e}\")
        raise
    logger.info(f\"Loaded {len(examples)} eval examples\")
    return examples

def generate_predictions(model, tokenizer, examples: List[Dict], max_new_tokens: int, batch_size: int) -> List[str]:
    \"\"\"Generate code predictions using Transformers 4.40 pipeline\"\"\"
    # Initialize text generation pipeline with FlashAttention if available
    try:
        gen_pipeline = pipeline(
            \"text-generation\",
            model=model,
            tokenizer=tokenizer,
            torch_dtype=torch.bfloat16,
            device_map=\"auto\",
            max_new_tokens=max_new_tokens,
            do_sample=False,
            num_return_sequences=1
        )
    except Exception as e:
        logger.error(f\"Failed to initialize pipeline: {e}\")
        raise

    predictions = []
    # Process in batches to avoid OOM
    for i in range(0, len(examples), batch_size):
        batch = examples[i:i+batch_size]
        prompts = [
            f\"<>You are an enterprise Java developer. Write code for the following task.<>[INST]{ex['input']}[/INST]\"
            for ex in batch
        ]
        try:
            batch_preds = gen_pipeline(prompts)
            for pred in batch_preds:
                # Extract generated code (remove prompt)
                generated = pred[0][\"generated_text\"]
                # Split on [/INST] to remove prompt part
                if \"[/INST]\" in generated:
                    code = generated.split(\"[/INST]\")[-1].strip()
                else:
                    code = generated.strip()
                predictions.append(code)
        except Exception as e:
            logger.error(f\"Generation failed for batch starting at {i}: {e}\")
            predictions.extend([\"\"] * len(batch))

    return predictions

def calculate_metrics(examples: List[Dict], predictions: List[str]) -> Dict:
    \"\"\"Calculate pass@1, edit similarity, and syntax validity\"\"\"
    # Calculate pass@1 using code_eval
    try:
        pass_at_1, _ = code_eval.compute(
            references=[ex[\"test_cases\"] for ex in examples],
            predictions=[[pred] for pred in predictions],
            k=[1]
        )
    except Exception as e:
        logger.error(f\"Failed to calculate pass@1: {e}\")
        pass_at_1 = {\"pass@1\": 0.0}

    # Calculate edit similarity
    try:
        edit_similarity = edit_sim.compute(
            references=[ex[\"expected_output\"] for ex in examples],
            predictions=predictions
        )
    except Exception as e:
        logger.error(f\"Failed to calculate edit similarity: {e}\")
        edit_similarity = {\"edit_similarity\": 0.0}

    # Calculate syntax validity (simple check for Java)
    valid_syntax = 0
    for pred in predictions:
        if pred.strip().startswith(\"public class\") or pred.strip().startswith(\"public interface\"):
            valid_syntax += 1
    syntax_pct = valid_syntax / len(predictions) if predictions else 0.0

    return {
        \"pass@1\": pass_at_1[\"pass@1\"],
        \"edit_similarity\": edit_similarity[\"edit_similarity\"],
        \"syntax_valid_pct\": syntax_pct,
        \"total_examples\": len(examples),
        \"total_predictions\": len(predictions)
    }

def main():
    args = parse_args()

    # Load model and tokenizer
    try:
        tokenizer = AutoTokenizer.from_pretrained(args.model_path, legacy=False)
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
        model = AutoModelForCausalLM.from_pretrained(
            args.model_path,
            torch_dtype=torch.bfloat16,
            device_map=\"auto\",
            attn_implementation=\"flash_attention_2\" if torch.cuda.is_available() else \"eager\"
        )
    except Exception as e:
        logger.error(f\"Failed to load model/tokenizer: {e}\")
        sys.exit(1)

    # Load eval data
    try:
        eval_examples = load_eval_data(args.eval_data)
    except Exception as e:
        logger.error(f\"Failed to load eval data: {e}\")
        sys.exit(1)

    # Generate predictions
    logger.info(\"Generating predictions...\")
    try:
        predictions = generate_predictions(
            model, tokenizer, eval_examples,
            args.max_new_tokens, args.batch_size
        )
    except Exception as e:
        logger.error(f\"Prediction generation failed: {e}\")
        sys.exit(1)

    # Calculate metrics
    logger.info(\"Calculating metrics...\")
    metrics = calculate_metrics(eval_examples, predictions)

    # Save results
    try:
        with open(args.output_path, \"w\") as f:
            json.dump(metrics, f, indent=2)
        logger.info(f\"Results saved to {args.output_path}\")
        logger.info(f\"Metrics: {json.dumps(metrics, indent=2)}\")
    except Exception as e:
        logger.error(f\"Failed to save results: {e}\")
        sys.exit(1)

if __name__ == \"__main__\":
    main()

import os
import sys
import json
import logging
from dataclasses import dataclass, field
from typing import Optional

import torch
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    TrainingArguments,
    Trainer,
    HfArgumentParser,
    set_seed
)
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training, TaskType
import datasets

# Logging config
logging.basicConfig(
    format=\"%(asctime)s - %(levelname)s - %(name)s - %(message)s\",
    level=logging.INFO,
    handlers=[logging.StreamHandler(sys.stdout)]
)
logger = logging.getLogger(__name__)

@dataclass
class LoRAFineTuneArgs:
    base_model_path: str = field(default=\"codellama/CodeLlama-7b-Instruct-hf\", metadata={\"help\": \"Base model for LoRA\"})
    train_data_path: str = field(default=\"enterprise_code_train.jsonl\", metadata={\"help\": \"Training data path\"})
    output_dir: str = field(default=\"./codellama-7b-lora-enterprise\", metadata={\"help\": \"Output directory\"})
    lora_r: int = field(default=16, metadata={\"help\": \"LoRA rank\"})
    lora_alpha: int = field(default=32, metadata={\"help\": \"LoRA alpha\"})
    lora_dropout: float = field(default=0.05, metadata={\"help\": \"LoRA dropout\"})
    max_seq_length: int = field(default=2048, metadata={\"help\": \"Max sequence length\"})
    use_4bit: bool = field(default=True, metadata={\"help\": \"Use 4-bit quantization for base model\"})

def load_dataset(data_path: str, tokenizer: AutoTokenizer, max_seq_length: int):
    \"\"\"Load and tokenize dataset for LoRA fine-tuning\"\"\"
    if not os.path.exists(data_path):
        raise FileNotFoundError(f\"Data not found at {data_path}\")

    examples = []
    try:
        with open(data_path, \"r\") as f:
            for line_num, line in enumerate(f, 1):
                try:
                    ex = json.loads(line.strip())
                    if \"input\" not in ex or \"output\" not in ex:
                        continue
                    examples.append(ex)
                except json.JSONDecodeError:
                    logger.warning(f\"Skipping invalid JSON on line {line_num}\")
                    continue
    except Exception as e:
        logger.error(f\"Failed to load dataset: {e}\")
        raise

    def tokenize_function(examples):
        prompts = [
            f\"<>Enterprise Java developer. Complete the code.<>[INST]{ex['input']}[/INST]\"
            for ex in examples
        ]
        full_texts = [p + ex[\"output\"] for p, ex in zip(prompts, examples)]
        tokenized = tokenizer(
            full_texts,
            max_length=max_seq_length,
            truncation=True,
            padding=\"max_length\",
            return_tensors=\"pt\"
        )
        # Create labels: ignore prompt tokens
        labels = tokenized[\"input_ids\"].clone()
        prompt_tokenized = tokenizer(prompts, return_tensors=\"pt\")[\"input_ids\"]
        for i in range(len(labels)):
            labels[i, :prompt_tokenized.shape[1]] = -100
        tokenized[\"labels\"] = labels
        return tokenized

    dataset = datasets.Dataset.from_list(examples)
    tokenized_dataset = dataset.map(
        lambda x: tokenize_function([x]),
        batched=False,
        remove_columns=dataset.column_names
    )
    return tokenized_dataset

def main():
    parser = HfArgumentParser((LoRAFineTuneArgs, TrainingArguments))
    if len(sys.argv) == 2 and sys.argv[1].endswith(\".json\"):
        lora_args, training_args = parser.parse_json_file(os.path.abspath(sys.argv[1]))
    else:
        lora_args, training_args = parser.parse_args_into_dataclasses()

    set_seed(training_args.seed)

    # Load tokenizer
    try:
        tokenizer = AutoTokenizer.from_pretrained(lora_args.base_model_path, legacy=False)
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
    except Exception as e:
        logger.error(f\"Tokenizer load failed: {e}\")
        sys.exit(1)

    # Load base model with 4-bit quantization if enabled
    try:
        if lora_args.use_4bit:
            from transformers import BitsAndBytesConfig
            bnb_config = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_use_double_quant=True,
                bnb_4bit_quant_type=\"nf4\",
                bnb_4bit_compute_dtype=torch.bfloat16
            )
            model = AutoModelForCausalLM.from_pretrained(
                lora_args.base_model_path,
                quantization_config=bnb_config,
                device_map=\"auto\",
                attn_implementation=\"flash_attention_2\"
            )
            model = prepare_model_for_kbit_training(model)
        else:
            model = AutoModelForCausalLM.from_pretrained(
                lora_args.base_model_path,
                torch_dtype=torch.bfloat16,
                device_map=\"auto\"
            )
    except Exception as e:
        logger.error(f\"Model load failed: {e}\")
        sys.exit(1)

    # Configure LoRA
    lora_config = LoraConfig(
        r=lora_args.lora_r,
        lora_alpha=lora_args.lora_alpha,
        lora_dropout=lora_args.lora_dropout,
        target_modules=[\"q_proj\", \"v_proj\", \"k_proj\", \"o_proj\"],  # CodeLlama attention modules
        task_type=TaskType.CAUSAL_LM,
        bias=\"none\"
    )
    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()  # Log trainable params count

    # Load and tokenize dataset
    try:
        tokenized_dataset = load_dataset(lora_args.train_data_path, tokenizer, lora_args.max_seq_length)
        train_test_split = tokenized_dataset.train_test_split(test_size=0.1)
        train_dataset = train_test_split[\"train\"]
        eval_dataset = train_test_split[\"test\"]
    except Exception as e:
        logger.error(f\"Dataset load failed: {e}\")
        sys.exit(1)

    # Initialize Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset
    )

    # Train
    logger.info(\"Starting LoRA fine-tuning...\")
    try:
        trainer.train()
        trainer.save_model(lora_args.output_dir)
        # Save LoRA adapter config
        model.save_pretrained(lora_args.output_dir)
        logger.info(f\"LoRA model saved to {lora_args.output_dir}\")
    except Exception as e:
        logger.error(f\"Training failed: {e}\")
        sys.exit(1)

if __name__ == \"__main__\":
    main()

# requirements.txt (always pin exact versions)
transformers==4.40.2
peft==0.7.1
datasets==2.19.1
accelerate==0.29.3

# Validate FlashAttention-2 compatibility
from transformers.utils import is_flash_attention_available
import torch

if is_flash_attention_available():
    try:
        # Test with max sequence length
        test_input = torch.randint(0, 32000, (1, 2048))
        model(test_input)  # Will throw error if incompatible
        print(\"FlashAttention-2 compatible\")
    except Exception as e:
        print(f\"FlashAttention-2 incompatible: {e}\")

# Add custom enterprise tokens to tokenizer
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained(\"codellama/CodeLlama-7b-Instruct-hf\")
custom_tokens = [\"BigDecimal\", \"Currency\", \"PaymentProcessor\", \"InternalApiV2\"]
num_added = tokenizer.add_tokens(custom_tokens)
print(f\"Added {num_added} custom tokens\")
# Resize model embeddings to match new tokenizer size
model.resize_token_embeddings(len(tokenizer))