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()