scripts/programming_finetune.py

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1"

import torch
from datasets import Dataset
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TrainingArguments,
    Trainer,
    DataCollatorForLanguageModeling,
    BitsAndBytesConfig,
)
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training, TaskType

MODEL_NAME = "Qwen/Qwen2.5-Coder-7B-Instruct"
OUTPUT_DIR = "/home/ai/qwen-coder-programming-finetuned"

PROG_DATA = [
    {"system": "You are an expert software engineer. Write clean, correct, efficient code.", "prompt": "Write a Python function to check if a string is a valid palindrome ignoring case and non-alphanumeric characters.", "output": "```python\ndef is_palindrome(s: str) -> bool:\n    cleaned = ''.join(c.lower() for c in s if c.isalnum())\n    return cleaned == cleaned[::-1]\n```\n\nTime: O(n), Space: O(n). Returns True for 'A man, a plan, a canal: Panama'."},
    {"system": "You are an expert software engineer. Write clean, correct, efficient code.", "prompt": "Write a Python function using a binary search to find the index of a target in a sorted list, or -1 if not present.", "output": "```python\ndef binary_search(arr, target):\n    low, high = 0, len(arr) - 1\n    while low <= high:\n        mid = (low + high) // 2\n        if arr[mid] == target:\n            return mid\n        elif arr[mid] < target:\n            low = mid + 1\n        else:\n            high = mid - 1\n    return -1\n```\n\nTime: O(log n), Space: O(1)." },
    {"system": "You are an expert software engineer. Write clean, correct, efficient code.", "prompt": "Write a Python function that returns the nth Fibonacci number using dynamic programming with memoization.", "output": "```python\nfrom functools import lru_cache\n\ndef fib(n: int) -> int:\n    @lru_cache(maxsize=None)\n    def _fib(n: int) -> int:\n        if n < 2:\n            return n\n        return _fib(n - 1) + _fib(n - 2)\n    return _fib(n)\n```\n\nTime: O(n), Space: O(n). Using lru_cache avoids exponential recomputation."},
    {"system": "You are an expert software engineer. Write clean, correct, efficient code.", "prompt": "Write a Python class implementing a simple LRU cache with get and put operations in O(1).", "output": "```python\nfrom collections import OrderedDict\n\nclass LRUCache:\n    def __init__(self, capacity: int):\n        self.capacity = capacity\n        self.cache = OrderedDict()\n\n    def get(self, key: int) -> int:\n        if key not in self.cache:\n            return -1\n        self.cache.move_to_end(key)\n        return self.cache[key]\n\n    def put(self, key: int, value: int) -> None:\n        if key in self.cache:\n            self.cache.move_to_end(key)\n        self.cache[key] = value\n        if len(self.cache) > self.capacity:\n            self.cache.popitem(last=False)\n```\n\nBoth operations are O(1) using OrderedDict."},
    {"system": "You are an expert software engineer. Write clean, correct, efficient code.", "prompt": "Write a Python function to merge two sorted lists into one sorted list.", "output": "```python\ndef merge_sorted(a, b):\n    result = []\n    i = j = 0\n    while i < len(a) and j < len(b):\n        if a[i] <= b[j]:\n            result.append(a[i]); i += 1\n        else:\n            result.append(b[j]); j += 1\n    result.extend(a[i:])\n    result.extend(b[j:])\n    return result\n```\n\nTime: O(n + m), Space: O(n + m)." },
    {"system": "You are an expert software engineer. Write clean, correct, efficient code.", "prompt": "Write a Python function using a trie to find if a word exists in a set of words.", "output": "```python\nclass TrieNode:\n    def __init__(self):\n        self.children = {}\n        self.is_word = False\n\ndef build_trie(words):\n    root = TrieNode()\n    for w in words:\n        node = root\n        for ch in w:\n            node = node.children.setdefault(ch, TrieNode())\n        node.is_word = True\n    return root\n\ndef exists(trie, word):\n    node = trie\n    for ch in word:\n        if ch not in node.children:\n            return False\n        node = node.children[ch]\n    return node.is_word\n```\n\nInsert and search are O(len(word))."},
    {"system": "You are an expert software engineer. Write clean, correct, efficient code.", "prompt": "Write a Python function that reverses a linked list in place.", "output": "```python\ndef reverse_linked_list(head):\n    prev = None\n    curr = head\n    while curr:\n        nxt = curr.next\n        curr.next = prev\n        prev = curr\n        curr = nxt\n    return prev\n```\n\nIterative approach, O(n) time, O(1) space."},
    {"system": "You are an expert software engineer. Write clean, correct, efficient code.", "prompt": "Write a Python function using a heap to return the k largest elements of a list.", "output": "```python\nimport heapq\n\ndef k_largest(nums, k):\n    if k <= 0:\n        return []\n    heap = nums[:k]\n    heapq.heapify(heap)\n    for x in nums[k:]:\n        if x > heap[0]:\n            heapq.heapreplace(heap, x)\n    return sorted(heap, reverse=True)\n```\n\nUses a min-heap of size k. Time: O(n log k)."},
    {"system": "You are an expert software engineer. Write clean, correct, efficient code.", "prompt": "Write a Python function that returns all permutations of a given list of distinct integers.", "output": "```python\ndef permutations(nums):\n    result = []\n    def backtrack(path, remaining):\n        if not remaining:\n            result.append(path[:])\n            return\n        for i in range(len(remaining)):\n            backtrack(path + [remaining[i]], remaining[:i] + remaining[i+1:])\n    backtrack([], nums)\n    return result\n```\n\nBacktracking, O(n!) permutations."},
    {"system": "You are an expert software engineer. Write clean, correct, efficient code.", "prompt": "Write Python code using dynamic programming to compute the length of the longest increasing subsequence in an array.", "output": "```python\ndef lis_length(nums):\n    import bisect\n    tails = []\n    for x in nums:\n        i = bisect.bisect_left(tails, x)\n        if i == len(tails):\n            tails.append(x)\n        else:\n            tails[i] = x\n    return len(tails)\n```\n\nPatience sorting: O(n log n) time, O(n) space."},
    {"system": "You are an expert software engineer. Write clean, correct, efficient code.", "prompt": "Write a Python function that checks whether a binary tree is balanced (heights differ by at most 1).", "output": "```python\ndef is_balanced(root):\n    def height(node):\n        if not node:\n            return 0\n        left = height(node.left)\n        right = height(node.right)\n        if left == -1 or right == -1 or abs(left - right) > 1:\n            return -1\n        return 1 + max(left, right)\n    return height(root) != -1\n```\n\nDFS returns -1 for unbalanced subtree. O(n) time."},
    {"system": "You are an expert software engineer. Write clean, correct, efficient code.", "prompt": "Write Python code using a queue (BFS) to find the shortest path length in a grid from (0,0) to (n-1,m-1) avoiding walls.", "output": "```python\nfrom collections import deque\n\ndef shortest_path(grid):\n    n, m = len(grid), len(grid[0])\n    if grid[0][0] == 1 or grid[n-1][m-1] == 1:\n        return -1\n    q = deque([(0, 0, 1)])\n    grid[0][0] = 1\n    while q:\n        r, c, d = q.popleft()\n        if (r, c) == (n-1, m-1):\n            return d\n        for dr, dc in ((1,0),(-1,0),(0,1),(0,-1)):\n            nr, nc = r+dr, c+dc\n            if 0 <= nr < n and 0 <= nc < m and grid[nr][nc] == 0:\n                grid[nr][nc] = 1\n                q.append((nr, nc, d+1))\n    return -1\n```\n\nBFS on grid: O(n*m)."},
    {"system": "You are an expert software engineer. Write clean, correct, efficient code.", "prompt": "Write Python code using a hash map to find two numbers in a list that sum to a target.", "output": "```python\ndef two_sum(nums, target):\n    seen = {}\n    for i, x in enumerate(nums):\n        complement = target - x\n        if complement in seen:\n            return [seen[complement], i]\n        seen[x] = i\n    return []\n```\n\nSingle pass, O(n) time, O(n) space."},
    {"system": "You are an expert software engineer. Write clean, correct, efficient code.", "prompt": "Write a Python function that converts an integer to its binary string representation without using bin().", "output": "```python\ndef to_binary(n: int) -> str:\n    if n == 0:\n        return \"0\"\n    bits = []\n    while n:\n        bits.append(str(n & 1))\n        n >>= 1\n    return ''.join(reversed(bits))\n```\n\nRepeatedly takes least significant bit. O(log n) time."},
    {"system": "You are an expert software engineer. Write clean, correct, efficient code.", "prompt": "Write Python code using Dijkstra's algorithm to find the shortest path in a weighted graph.", "output": "```python\nimport heapq\n\ndef dijkstra(graph, start):\n    dist = {node: float('inf') for node in graph}\n    dist[start] = 0\n    pq = [(0, start)]\n    while pq:\n        d, u = heapq.heappop(pq)\n        if d > dist[u]:\n            continue\n        for v, w in graph[u]:\n            nd = d + w\n            if nd < dist[v]:\n                dist[v] = nd\n                heapq.heappush(pq, (nd, v))\n    return dist\n```\n\nPriority-queue Dijkstra: O((V+E) log V)."},
    {"system": "You are an expert software engineer. Write clean, correct, efficient code.", "prompt": "Write a Python function to determine if two strings are anagrams.", "output": "```python\nfrom collections import Counter\n\ndef is_anagram(s: str, t: str) -> bool:\n    if len(s) != len(t):\n        return False\n    return Counter(s) == Counter(t)\n```\n\nCounter comparison is O(n). Works for any character set."},
    {"system": "You are an expert software engineer. Write clean, correct, efficient code.", "prompt": "Write Python code with a DFS to find all paths from a source to a target in an undirected graph.", "output": "```python\ndef all_paths(graph, src, dst):\n    result = []\n    def dfs(node, path):\n        path.append(node)\n        if node == dst:\n            result.append(path[:])\n        else:\n            for nei in graph[node]:\n                if nei not in path:\n                    dfs(nei, path)\n        path.pop()\n    dfs(src, [])\n    return result\n```\n\nBacktracking DFS enumerates all paths."},
]

def format_example(ex):
    return f"""<|im_start|>system
{ex['system']}<|im_end|>
<|im_start|>user
{ex['prompt']}<|im_end|>
<|im_start|>assistant
{ex['output']}<|im_end|>"""

def main():
    print("Loading tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(
        MODEL_NAME, trust_remote_code=True, use_fast=True
    )
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"

    print("Loading model with 4-bit quantization...")
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16,
        bnb_4bit_use_double_quant=True,
    )

    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        quantization_config=bnb_config,
        device_map="auto",
        trust_remote_code=True,
        torch_dtype=torch.bfloat16,
    )

    model = prepare_model_for_kbit_training(model)
    model.gradient_checkpointing_enable()

    lora_config = LoraConfig(
        r=32,
        lora_alpha=64,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
        lora_dropout=0.05,
        bias="none",
        task_type=TaskType.CAUSAL_LM,
    )

    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()

    print("Preparing dataset...")
    formatted_data = [{"text": format_example(ex)} for ex in PROG_DATA]
    dataset = Dataset.from_list(formatted_data)

    def tokenize_function(examples):
        result = tokenizer(
            examples["text"],
            truncation=True,
            max_length=1024,
            padding="max_length",
            return_tensors="pt",
        )
        result["labels"] = result["input_ids"].clone()
        return result

    tokenized_dataset = dataset.map(
        tokenize_function, batched=True, remove_columns=["text"]
    )
    tokenized_dataset.set_format(
        type="torch", columns=["input_ids", "attention_mask", "labels"]
    )

    training_args = TrainingArguments(
        output_dir=OUTPUT_DIR,
        num_train_epochs=2,
        per_device_train_batch_size=2,
        gradient_accumulation_steps=4,
        warmup_steps=5,
        learning_rate=2e-4,
        fp16=False,
        bf16=True,
        logging_steps=5,
        save_steps=500,
        save_total_limit=2,
        optim="paged_adamw_8bit",
        lr_scheduler_type="cosine",
        report_to="none",
        remove_unused_columns=False,
        dataloader_pin_memory=False,
        max_grad_norm=0.3,
    )

    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

    trainer = Trainer(
        model=model,
        train_dataset=tokenized_dataset,
        args=training_args,
        data_collator=data_collator,
    )

    print("Starting training...")
    trainer.train()

    print("Saving model...")
    trainer.save_model()
    tokenizer.save_pretrained(OUTPUT_DIR)
    print(f"Model saved to {OUTPUT_DIR}")

if __name__ == "__main__":
    main()