from transformers import AutoModelForCausalLM, AutoTokenizer model = AutoModelForCausalLM.from_pretrained("distil-labs/distil-qwen3-4b-text2sql")tokenizer = AutoTokenizer.from_pretrained("distil-labs/distil-qwen3-4b-text2sql") schema = """CREATE TABLE employees ( id INTEGER PRIMARY KEY, name TEXT NOT NULL, department TEXT, salary INTEGER);""" question = "How many employees earn more than 50000?" messages = [ { "role": "system", "content": """You are a problem solving model working on task_description XML block:<task_description>You are given a database schema and a natural language question. Generate the SQL query that answers the question. Input:- Schema: One or two table definitions in SQL DDL format- Question: Natural language question about the data Output:- A single SQL query that answers the question- No explanations, comments, or additional text Rules:- Use only tables and columns from the provided schema- Use uppercase SQL keywords (SELECT, FROM, WHERE, etc.)- Use SQLite-compatible syntax</task_description>You will be given a single task in the question XML blockSolve only the task in question block.Generate only the answer, do not generate anything else""" }, { "role": "user", "content": f"""Now for the real task, solve the task in question block.Generate only the solution, do not generate anything else<question>Schema:{schema} Question: {question}</question>""" }] text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)inputs = tokenizer(text, return_tensors="pt")outputs = model.generate(**inputs, max_new_tokens=256, temperature=0)print(tokenizer.decode(outputs[0], skip_special_tokens=True))