使用chatdb/natural-sql-7b构建Text-to-SQL遇无输出及CUDA报错求助

一、CUDA不可用问题解决 #

• 重新安装带CUDA支持的PyTorch：根据你的CUDA版本（如cu118）执行对应安装命令，例如：

  pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
  

• 验证CUDA可用性：运行以下代码确认返回True：

  import torch
  print(torch.cuda.is_available())
  

• 优化模型加载策略：若CUDA可用，device_map="auto"会自动将模型加载到GPU；若显存不足，可开启4bit量化（需先安装bitsandbytes库），修改模型加载代码：

  model = AutoModelForCausalLM.from_pretrained(
      "chatdb/natural-sql-7b",
      device_map="auto",
      torch_dtype=torch.float16,
      load_in_4bit=True  # 开启4bit量化节省显存
  )
  

二、模型卡住不生成结果问题解决 #

1. 修正Prompt模板 #

原Prompt中Answer部分的冗余内容会干扰模型生成逻辑，修改为：

prompt = f"""
### Task 

Generate a SQL query to answer the following question: `{question}` 

### PostgreSQL Database Schema 
The query will run on a database with the following schema: 

{get_schema()}

### Answer 
Here is the SQL query that answers the question:
```sql
"""

2. 调整生成参数 #

• 替换硬编码的结束符ID为tokenizer内置值，避免匹配错误
• 添加early_stopping=True确保模型生成结束符后停止
• 将输入张量移至对应设备（GPU优先）

修改后的生成代码：

device = "cuda" if torch.cuda.is_available() else "cpu"
inputs = tokenizer(prompt, return_tensors="pt").to(device)

generated_ids = model.generate(
    **inputs,
    num_return_sequences=1,
    eos_token_id=tokenizer.eos_token_id,
    pad_token_id=tokenizer.pad_token_id,
    max_new_tokens=400,
    do_sample=False,
    num_beams=1,
    early_stopping=True
)

3. 避免CPU运行瓶颈 #

7B模型在CPU上运行速度极慢，会表现为“卡住”，优先使用GPU；若必须用CPU，需确保设备有足够内存（至少16GB以上）。

修改后的完整代码示例 #

import torch
from db import get_schema
from transformers import AutoModelForCausalLM, AutoTokenizer

# 自动选择运行设备
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")

tokenizer = AutoTokenizer.from_pretrained("chatdb/natural-sql-7b")
model = AutoModelForCausalLM.from_pretrained(
    "chatdb/natural-sql-7b",
    device_map="auto",
    torch_dtype=torch.float16,
    load_in_4bit=True  # 显存不足时开启
)

question = 'How many employees are there?'

prompt = f"""
### Task 

Generate a SQL query to answer the following question: `{question}` 

### PostgreSQL Database Schema 
The query will run on a database with the following schema: 

{get_schema()}

### Answer 
Here is the SQL query that answers the question:
```sql
"""

print("Question: " + question)
print("SQL: ")

inputs = tokenizer(prompt, return_tensors="pt").to(device)

generated_ids = model.generate(
    **inputs,
    num_return_sequences=1,
    eos_token_id=tokenizer.eos_token_id,
    pad_token_id=tokenizer.pad_token_id,
    max_new_tokens=400,
    do_sample=False,
    num_beams=1,
    early_stopping=True
)

outputs = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
print(outputs[0].split("```sql")[-1].strip())

内容的提问来源于stack exchange，提问作者Abhra Sarkar