update endpoint helper files

finetune_coding_llm_colab.py

@@ -0,0 +1,246 @@
+"""
+Local LoRA fine-tuning script for a small coding model.
+
+Quick start (Windows/Linux local):
+1) pip install transformers datasets peft accelerate bitsandbytes huggingface_hub
+2) python finetune_coding_llm_colab.py --dataset-size 8000
+3) Optional upload:
+   python finetune_coding_llm_colab.py --skip-train --upload --hf-repo your-user/your-model
+"""
+
+import argparse
+import json
+import os
+import random
+import torch
+from datasets import load_dataset
+from huggingface_hub import upload_folder
+from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    BitsAndBytesConfig,
+    Trainer,
+    TrainingArguments,
+)
+
+
+DEFAULT_MODEL_NAME = "Qwen/Qwen2.5-Coder-0.5B-Instruct"
+DEFAULT_OUTPUT_DIR = "./model"
+DEFAULT_TRAIN_FILE = "train.json"
+HF_REPO_ID = "your-username/coding-llm-model"
+
+# Keep dataset size in the requested 5k-10k window.
+DATASET_SIZE = 8000
+
+
+TEMPLATES = [
+    {
+        "instruction": "Fix the Python code",
+        "input": "def add(a,b) return a+b",
+        "output": "def add(a, b): return a + b",
+        "explanation": "Added missing colon and corrected syntax.",
+    },
+    {
+        "instruction": "Fix loop syntax",
+        "input": "for i in range(5 print(i)",
+        "output": "for i in range(5): print(i)",
+        "explanation": "Added missing parenthesis and colon.",
+    },
+    {
+        "instruction": "Fix condition",
+        "input": "if x = 10: print(x)",
+        "output": "if x == 10: print(x)",
+        "explanation": "Corrected assignment to comparison operator.",
+    },
+    {
+        "instruction": "Explain code",
+        "input": "for i in range(3): print(i)",
+        "output": "Prints numbers from 0 to 2.",
+        "explanation": "Loop iterates from 0 to 2 and prints values.",
+    },
+]
+
+
+def format_training_text(template):
+    target = {
+        "code": template["output"],
+        "explanation": template["explanation"],
+    }
+    return (
+        f"Instruction: {template['instruction']}\n"
+        f"Input: {template['input']}\n"
+        "Return only valid JSON with keys code and explanation.\n"
+        f"JSON: {json.dumps(target, ensure_ascii=False)}\n"
+    )
+
+
+def generate_sample():
+    template = random.choice(TEMPLATES)
+    text = format_training_text(template)
+    return {
+        "instruction": template["instruction"],
+        "input": template["input"],
+        "output": template["output"],
+        "explanation": template["explanation"],
+        "text": text,
+        "confidence": round(random.uniform(0.9, 0.99), 2),
+        "relevancy": round(random.uniform(0.85, 0.99), 2),
+    }
+
+
+def build_dataset(train_file, size=DATASET_SIZE):
+    dataset = [generate_sample() for _ in range(size)]
+    with open(train_file, "w", encoding="utf-8") as f:
+        json.dump(dataset, f, indent=2)
+    print(f"Dataset created: {len(dataset)} samples -> {train_file}")
+
+
+def run_training(
+    model_name,
+    train_file,
+    output_dir,
+    epochs,
+    batch_size,
+    learning_rate,
+    max_length,
+    max_train_samples,
+    use_4bit,
+):
+    if not os.path.exists(train_file):
+        raise FileNotFoundError(
+            f"Training file not found: {train_file}. Generate it with generate_dataset.py first."
+        )
+
+    dataset = load_dataset("json", data_files=train_file)
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    tokenizer.pad_token = tokenizer.eos_token
+
+    def format_data(example):
+        text = example.get("text")
+        if not text:
+            text = format_training_text(example)
+        tokens = tokenizer(
+            text,
+            truncation=True,
+            padding="max_length",
+            max_length=max_length,
+        )
+        tokens["labels"] = tokens["input_ids"].copy()
+        return tokens
+
+    tokenized = dataset.map(
+        format_data,
+        remove_columns=dataset["train"].column_names,
+        desc="Tokenizing training dataset",
+    )
+    if max_train_samples > 0:
+        max_train_samples = min(max_train_samples, len(tokenized["train"]))
+        tokenized["train"] = tokenized["train"].select(range(max_train_samples))
+
+    fp16_enabled = torch.cuda.is_available()
+    quantize_4bit = use_4bit and torch.cuda.is_available()
+    if use_4bit and not torch.cuda.is_available():
+        print("Warning: --use-4bit requested but CUDA not available. Falling back to standard loading.")
+    if quantize_4bit:
+        bnb_config = BitsAndBytesConfig(load_in_4bit=True)
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            quantization_config=bnb_config,
+            device_map="auto",
+        )
+        model = prepare_model_for_kbit_training(model)
+    else:
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            device_map="auto" if torch.cuda.is_available() else None,
+        )
+
+    lora_config = LoraConfig(
+        r=8,
+        lora_alpha=16,
+        target_modules=["q_proj", "v_proj"],
+        lora_dropout=0.05,
+        bias="none",
+        task_type="CAUSAL_LM",
+    )
+    model = get_peft_model(model, lora_config)
+
+    training_args = TrainingArguments(
+        output_dir=output_dir,
+        per_device_train_batch_size=batch_size,
+        num_train_epochs=epochs,
+        gradient_accumulation_steps=2,
+        logging_steps=10,
+        save_steps=100,
+        learning_rate=learning_rate,
+        fp16=fp16_enabled,
+        dataloader_pin_memory=torch.cuda.is_available(),
+        report_to="none",
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=tokenized["train"],
+    )
+    trainer.train()
+
+    model.save_pretrained(output_dir)
+    tokenizer.save_pretrained(output_dir)
+    print(f"Model and tokenizer saved to: {output_dir}")
+
+
+def upload_to_hf(repo_id, output_dir):
+    if not os.path.exists(output_dir):
+        raise FileNotFoundError(
+            f"Model output folder not found: {output_dir}. Run training before upload."
+        )
+
+    upload_folder(
+        folder_path=output_dir,
+        repo_id=repo_id,
+        repo_type="model",
+    )
+    print(f"Uploaded to Hugging Face repo: {repo_id}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dataset-size", type=int, default=DATASET_SIZE)
+    parser.add_argument("--train-file", type=str, default=DEFAULT_TRAIN_FILE)
+    parser.add_argument("--output-dir", type=str, default=DEFAULT_OUTPUT_DIR)
+    parser.add_argument("--model-name", type=str, default=DEFAULT_MODEL_NAME)
+    parser.add_argument("--epochs", type=float, default=1)
+    parser.add_argument("--batch-size", type=int, default=2)
+    parser.add_argument("--learning-rate", type=float, default=2e-4)
+    parser.add_argument("--max-length", type=int, default=512)
+    parser.add_argument("--max-train-samples", type=int, default=0)
+    parser.add_argument("--use-4bit", action="store_true")
+    parser.add_argument("--skip-dataset-gen", action="store_true")
+    parser.add_argument("--skip-train", action="store_true")
+    parser.add_argument("--upload", action="store_true")
+    parser.add_argument("--hf-repo", type=str, default=HF_REPO_ID)
+    args = parser.parse_args()
+
+    if not (5000 <= args.dataset_size <= 10000):
+        raise ValueError("dataset-size must be between 5000 and 10000")
+
+    if not args.skip_dataset_gen:
+        build_dataset(train_file=args.train_file, size=args.dataset_size)
+
+    if not args.skip_train:
+        run_training(
+            model_name=args.model_name,
+            train_file=args.train_file,
+            output_dir=args.output_dir,
+            epochs=args.epochs,
+            batch_size=args.batch_size,
+            learning_rate=args.learning_rate,
+            max_length=args.max_length,
+            max_train_samples=args.max_train_samples,
+            use_4bit=args.use_4bit,
+        )
+
+    if args.upload:
+        upload_to_hf(repo_id=args.hf_repo, output_dir=args.output_dir)