File size: 9,633 Bytes
f238825 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 | import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
import torch
from datasets import load_dataset, Dataset
from transformers import (
AutoModelForCausalLM,
AutoTokenizer,
TrainingArguments,
Trainer,
BitsAndBytesConfig,
)
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training, TaskType
from dataclasses import dataclass
from transformers import DataCollatorWithPadding
MODEL_NAME = "Qwen/Qwen2.5-Coder-7B-Instruct"
OUTPUT_DIR = "/home/ai/qwen-coder-programming-best"
SYSTEM = "You are an expert software engineer. Write clean, correct, efficient, well-structured code with brief explanations."
EXPERT_EXAMPLES = [
{"instruction": "Write a Python function to check if a string is a valid palindrome ignoring case and non-alphanumeric characters.", "input": "", "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'."},
{"instruction": "Write a Python function using a binary search to find the index of a target in a sorted list, or -1 if not present.", "input": "", "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)."},
{"instruction": "Write a Python class implementing a simple LRU cache with get and put operations in O(1).", "input": "", "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."},
{"instruction": "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.", "input": "", "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)."},
{"instruction": "Write Python code using Dijkstra's algorithm to find the shortest path in a weighted graph.", "input": "", "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((E+V) log V)."},
{"instruction": "Write a Python function using dynamic programming to compute the length of the longest increasing subsequence in an array.", "input": "", "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."},
]
# ----------------------------------------------------------------
# Load and filter CodeAlpaca-18k
ds = load_dataset("iamtarun/python_code_instructions_18k_alpaca", split="train")
def filter_quality(ex):
out = ex["output"] or ""
ins = ex["instruction"] or ""
if not ins.strip() or len(ins) > 2000:
return False
if len(out) < 10 or len(out) > 2000:
return False
if "def " not in out and "import " not in out and "class " not in out and "{" not in out and "function" not in out.lower():
return False
return True
ds = ds.filter(filter_quality, num_proc=4)
ds = ds.shuffle(seed=42)
N = 6000
ds = ds.select(range(min(N, len(ds))))
expert_ds = Dataset.from_list(EXPERT_EXAMPLES)
def normalize(ex):
return {"instruction": ex["instruction"], "input_text": str(ex.get("input") or ""), "output": ex["output"]}
ds = ds.map(normalize, remove_columns=ds.column_names)
expert_ds = expert_ds.map(lambda e: {"instruction": e["instruction"], "input_text": str(e.get("input") or ""), "output": e["output"]})
ds = Dataset.from_list(expert_ds.to_list() + ds.to_list())
print(f"Final train examples: {len(ds)}")
# ----------------------------------------------------------------
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"
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()
model.enable_input_require_grads()
lora_config = LoraConfig(
r=64,
lora_alpha=128,
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,
use_rslora=True,
)
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
def format_chat(ex):
user_content = ex["instruction"] + (f"\n\nInput:\n{ex['input_text']}" if ex["input_text"] else "")
prefix = f"<|im_start|>system\n{SYSTEM}<|im_end|>\n<|im_start|>user\n{user_content}<|im_end|>\n"
suffix = f"<|im_start|>assistant\n{ex['output']}<|im_end|>"
return {"prefix": prefix, "suffix": suffix}
def tokenize_with_masks(ex, max_length=2048):
p = tokenizer(ex["prefix"], add_special_tokens=False)
s = tokenizer(ex["suffix"], add_special_tokens=False)
input_ids = (p["input_ids"] + s["input_ids"])[:max_length]
labels = ([-100] * len(p["input_ids"]) + s["input_ids"])[:max_length]
attention_mask = [1] * len(input_ids)
return {"input_ids": input_ids, "labels": labels, "attention_mask": attention_mask}
ds_map = ds.map(format_chat, remove_columns=ds.column_names)
ds_map = ds_map.map(
lambda e: tokenize_with_masks(e),
remove_columns=ds_map.column_names,
batched=False,
)
print("Sample input ids len:", len(ds_map[0]["input_ids"]))
class CompletionOnlyDataCollator:
def __init__(self, tokenizer):
self.tokenizer = tokenizer
def __call__(self, features):
pad_id = self.tokenizer.pad_token_id
max_len = max(len(f["input_ids"]) for f in features)
input_ids, attention_mask, labels = [], [], []
for f in features:
ids, mask, lab = f["input_ids"], f["attention_mask"], f["labels"]
pad = max_len - len(ids)
input_ids.append(ids + [pad_id] * pad)
attention_mask.append(mask + [0] * pad)
labels.append(lab + [-100] * pad)
return {
"input_ids": torch.tensor(input_ids, dtype=torch.long),
"attention_mask": torch.tensor(attention_mask, dtype=torch.long),
"labels": torch.tensor(labels, dtype=torch.long),
}
training_args = TrainingArguments(
output_dir=OUTPUT_DIR,
num_train_epochs=3,
per_device_train_batch_size=8,
gradient_accumulation_steps=4,
warmup_steps=80,
learning_rate=2e-4,
lr_scheduler_type="cosine",
weight_decay=0.0,
fp16=False,
bf16=True,
max_grad_norm=1.0,
logging_steps=20,
save_steps=250,
save_total_limit=2,
optim="adamw_8bit",
report_to="none",
remove_unused_columns=False,
dataloader_pin_memory=False,
gradient_checkpointing_kwargs={"use_reentrant": False},
)
collator = CompletionOnlyDataCollator(tokenizer=tokenizer)
trainer = Trainer(
model=model,
train_dataset=ds_map,
args=training_args,
data_collator=collator,
)
print("Starting aggressive training...")
trainer.train()
trainer.save_model()
tokenizer.save_pretrained(OUTPUT_DIR)
print(f"Model saved to {OUTPUT_DIR}")
if __name__ == "__main__":
main() |