Spaces:
Running
Running
File size: 4,314 Bytes
a39d8ef | 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 | import os
# CRITICAL: Ye line sabse upar honi chahiye kisi bhi PyTorch import se pehle!
os.environ["CUDA_VISIBLE_DEVICES"] = "0,7"
import sys
import torch
from datasets import Dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig
from trl import GRPOConfig, GRPOTrainer
sys.path.insert(0, "./server")
from environment import NL2SQLEnvironment
from models import NL2SQLAction
from tasks import all_task_names, get_task
MODEL_NAME = "Qwen/Qwen2.5-Coder-7B-Instruct"
OUTPUT_DIR = "./qwen-7b-coder-nl2sql-grpo"
SYSTEM_PROMPT = """You are a Senior Database Architect and an expert in SQLite.
Your task is to translate natural language questions into highly optimized, correct SQLite SELECT queries.
STRICT RULES:
1. Output EXACTLY ONE valid SQLite query.
2. DO NOT wrap the query in markdown formatting (no ```sql or ```).
3. DO NOT output any explanations, conversational text, or preambles (e.g., never say "Here is the query").
4. ONLY use standard SQLite functions. Avoid SQL Server, MySQL, or PostgreSQL specific syntax.
5. If the question implies ordering, use the correct ORDER BY clause.
Your output must be executable directly against the database as-is."""
def build_dataset():
data = []
for t_name in all_task_names():
task = get_task(t_name)
schema = task.schema_context()
for ex in task.examples:
user_content = f"SCHEMA:\n{schema}\n\nQUESTION: {ex.question}"
data.append({
"prompt": [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": user_content}
],
"task_name": t_name
})
return Dataset.from_list(data)
def sql_reward_func(prompts, completions, task_name, **kwargs):
rewards = []
env = NL2SQLEnvironment()
for idx, completion in enumerate(completions):
generated_text = completion[0]['content'] if isinstance(completion, list) else completion
if generated_text.startswith("```"):
lines = generated_text.split("\n")
generated_text = "\n".join(l for l in lines if not l.strip().startswith("```")).strip()
current_task = task_name[idx] if isinstance(task_name, list) else task_name
env.reset(task_name=current_task)
try:
action = NL2SQLAction(query=generated_text)
obs = env.step(action)
rewards.append(float(obs.reward))
except Exception:
rewards.append(0.0)
return rewards
def main():
dataset = build_dataset()
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, padding_side="right")
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
model = AutoModelForCausalLM.from_pretrained(
MODEL_NAME,
torch_dtype=torch.bfloat16,
attn_implementation="sdpa" # Defaulting to sdpa to avoid any flash_attn setup issues
)
peft_config = LoraConfig(
r=128,
lora_alpha=256,
target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
bias="none",
task_type="CAUSAL_LM"
)
training_args = GRPOConfig(
output_dir=OUTPUT_DIR,
learning_rate=2e-5,
per_device_train_batch_size=2,
gradient_accumulation_steps=4,
max_completion_length=256,
num_generations=8,
temperature=0.5,
bf16=True,
logging_steps=5,
num_train_epochs=10,
report_to="none",
remove_unused_columns=False,
ddp_find_unused_parameters=False
)
trainer = GRPOTrainer(
model=model,
reward_funcs=sql_reward_func,
args=training_args,
train_dataset=dataset,
peft_config=peft_config,
processing_class=tokenizer
)
trainer.train()
if trainer.accelerator.is_main_process:
trainer.model.save_pretrained(f"{OUTPUT_DIR}/final")
tokenizer.save_pretrained(f"{OUTPUT_DIR}/final")
if __name__ == "__main__":
main() |