440 lines
16 KiB
Python
440 lines
16 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
将 Example_text2sql.md 解析为结构化数据集
|
||
|
||
输出格式：
|
||
1. JSONL格式（每行一个样本）
|
||
2. 按难度/类别分类
|
||
3. 包含Schema、SQL、执行结果、评分
|
||
"""
|
||
|
||
import re
|
||
import json
|
||
import yaml
|
||
from pathlib import Path
|
||
from typing import List, Dict, Optional, Any
|
||
from dataclasses import dataclass, asdict
|
||
from datetime import datetime
|
||
|
||
@dataclass
|
||
class Text2SQLSample:
|
||
"""单条Text2SQL样本"""
|
||
qid: str # 问题ID，如 "Q1"
|
||
question_zh: str # 中文问题
|
||
sql: str # SQL语句
|
||
schema_info: Dict[str, Any] # Schema信息（表、字段）
|
||
explanation: str = "" # SQL理由
|
||
question_en: Optional[str] = None # 英文问题（如有）
|
||
rating: Optional[int] = None # 准确度评分 1-10
|
||
execution_result: Optional[List[Dict]] = None # 执行结果（表格）
|
||
result_note: str = "" # 结果说明（如"当前查询无结果"）
|
||
tags: List[str] = None # 标签（如["聚合", "JOIN", "日期筛选"]）
|
||
difficulty: str = "medium" # 难度：easy/medium/hard
|
||
|
||
def __post_init__(self):
|
||
if self.tags is None:
|
||
self.tags = []
|
||
|
||
def to_dict(self):
|
||
return asdict(self)
|
||
|
||
def to_fewshot_example(self) -> str:
|
||
"""转换为few-shot示例字符串"""
|
||
return f"""问题：{self.question_zh}
|
||
Schema：{json.dumps(self.schema_info, ensure_ascii=False)[:500]}
|
||
SQL：
|
||
{self.sql}
|
||
"""
|
||
|
||
|
||
class ExampleText2SQLParser:
|
||
"""解析 Example_text2sql.md 文件"""
|
||
|
||
def __init__(self, md_path: str):
|
||
self.md_path = Path(md_path)
|
||
self.content = self.md_path.read_text(encoding='utf-8')
|
||
self.samples: List[Text2SQLSample] = []
|
||
|
||
def parse(self) -> List[Text2SQLSample]:
|
||
"""解析整个MD文件"""
|
||
# 按问题分割（## Qxx. 标题）
|
||
q_blocks = re.split(r'\n## Q(\d+)\. ', self.content)
|
||
|
||
# q_blocks[0] 是文件头（目录、固定日期说明等），跳过
|
||
for i in range(1, len(q_blocks)):
|
||
qid_num = q_blocks[i].split('\n')[0].strip() # 从分割中获取问题号
|
||
# 实际qid应为 "Q{num}"
|
||
# 重新用正则匹配完整的Q标题
|
||
pass
|
||
|
||
# 更稳健的方法：直接用正则查找所有Q标题块
|
||
pattern = r'(## Q(\d+)\. .*?)(?=\n## Q\d+\. |\Z)'
|
||
matches = re.findall(pattern, self.content, re.DOTALL)
|
||
|
||
print(f"找到 {len(matches)} 个问答对")
|
||
|
||
for full_match, qid in matches:
|
||
sample = self._parse_single_block(qid, full_match)
|
||
if sample:
|
||
self.samples.append(sample)
|
||
|
||
return self.samples
|
||
|
||
def _parse_single_block(self, qid: str, block: str) -> Optional[Text2SQLSample]:
|
||
"""解析单个问答块"""
|
||
try:
|
||
lines = block.strip().split('\n')
|
||
|
||
# 提取问题（支持中英文双语）
|
||
# 格式：### 问题 后跟英文段落，空行，然后中文段落
|
||
question_zh = ""
|
||
question_en = ""
|
||
found_question_header = False
|
||
paragraph_lines = []
|
||
paragraphs = []
|
||
|
||
for line in lines:
|
||
if line.startswith('### 问题'):
|
||
found_question_header = True
|
||
continue
|
||
if not found_question_header:
|
||
continue
|
||
if line.startswith('###'):
|
||
# 遇到下一节，停止
|
||
break
|
||
if line.strip() == '':
|
||
# 空行：段落结束
|
||
if paragraph_lines:
|
||
paragraphs.append(' '.join(paragraph_lines).strip())
|
||
paragraph_lines = []
|
||
else:
|
||
paragraph_lines.append(line.strip())
|
||
|
||
# 最后一个段落
|
||
if paragraph_lines:
|
||
paragraphs.append(' '.join(paragraph_lines).strip())
|
||
|
||
# 第一段为英文，第二段为中文
|
||
if len(paragraphs) >= 1:
|
||
question_en = paragraphs[0]
|
||
if len(paragraphs) >= 2:
|
||
question_zh = paragraphs[1]
|
||
|
||
# 提取SQL
|
||
sql_match = re.search(r'```sql\n(.*?)\n```', block, re.DOTALL)
|
||
if not sql_match:
|
||
return None
|
||
sql = sql_match.group(1).strip()
|
||
|
||
# 提取SQL理由（### SQL理由 后面的内容）
|
||
reason_match = re.search(r'### SQL理由\n\n(.*?)(?=\n### |\Z)', block, re.DOTALL)
|
||
explanation = reason_match.group(1).strip() if reason_match else ""
|
||
|
||
# 提取评分（### SQL准确度评分）
|
||
rating_match = re.search(r'\*\*评分:\s*(\d+)/10\*\*', block)
|
||
rating = int(rating_match.group(1)) if rating_match else None
|
||
|
||
# 提取结果说明（### SQL执行结果 后面）
|
||
result_match = re.search(r'### SQL执行结果\n\n(.*?)(?=\n### |\Z)', block, re.DOTALL)
|
||
result_note = result_match.group(1).strip() if result_match else ""
|
||
|
||
# 推断标签
|
||
tags = self._infer_tags(question_zh, sql)
|
||
|
||
# 推断难度
|
||
difficulty = self._infer_difficulty(sql, tags)
|
||
|
||
# 构建Schema信息（从SQL中提取的表和字段）
|
||
schema_info = self._extract_schema_from_sql(sql)
|
||
|
||
return Text2SQLSample(
|
||
qid=f"Q{qid}",
|
||
question_zh=question_zh.strip(),
|
||
question_en=question_en.strip() if question_en else None,
|
||
sql=sql,
|
||
schema_info=schema_info,
|
||
explanation=explanation,
|
||
rating=rating,
|
||
result_note=result_note,
|
||
tags=tags,
|
||
difficulty=difficulty
|
||
)
|
||
|
||
except Exception as e:
|
||
print(f"解析 {qid} 失败: {e}")
|
||
return None
|
||
|
||
def _infer_tags(self, question: str, sql: str) -> List[str]:
|
||
"""推断问题标签"""
|
||
tags = []
|
||
q_lower = question.lower()
|
||
sql_upper = sql.upper()
|
||
|
||
# 按业务域
|
||
if any(k in q_lower for k in ['对手方', 'broker', '经纪商']):
|
||
tags.append('broker')
|
||
if any(k in q_lower for k in ['账户', 'account', '客户']):
|
||
tags.append('account')
|
||
if any(k in q_lower for k in ['现金', 'cash', '余额']):
|
||
tags.append('cash')
|
||
if any(k in q_lower for k in ['持仓', 'holding', 'position']):
|
||
tags.append('holding')
|
||
if any(k in q_lower for k in ['交易', 'trade', 'commission']):
|
||
tags.append('trade')
|
||
if any(k in q_lower for k in ['IPO', '认购', 'subscription']):
|
||
tags.append('ipo')
|
||
if any(k in q_lower for k in ['公司行动', 'corporate action', '权益']):
|
||
tags.append('corporate_action')
|
||
if any(k in q_lower for k in ['汇率', 'exchange rate']):
|
||
tags.append('exchange_rate')
|
||
if any(k in q_lower for k in ['保证金', 'margin']):
|
||
tags.append('margin')
|
||
if any(k in q_lower for k in ['利息', 'interest']):
|
||
tags.append('interest')
|
||
if any(k in q_lower for k in ['报表', 'statement']):
|
||
tags.append('statement')
|
||
|
||
# 按SQL特征
|
||
if 'GROUP BY' in sql_upper:
|
||
tags.append('aggregation')
|
||
if 'JOIN' in sql_upper:
|
||
tags.append('join')
|
||
if 'WHERE' in sql_upper:
|
||
tags.append('filter')
|
||
if re.search(r'DATE.*BETWEEN|>=.*AND', sql_upper) or 'BETWEEN' in sql_upper:
|
||
tags.append('date_range')
|
||
if re.search(r'\bCOUNT\b|\bSUM\b|\bAVG\b|\bMAX\b|\bMIN\b', sql_upper):
|
||
tags.append('aggregation')
|
||
if 'ORDER BY' in sql_upper:
|
||
tags.append('ordering')
|
||
if 'TOP' in sql_upper or 'LIMIT' in sql_upper:
|
||
tags.append('limit')
|
||
if 'LEFT JOIN' in sql_upper:
|
||
tags.append('left_join')
|
||
if 'UNION' in sql_upper:
|
||
tags.append('union')
|
||
if re.search(r'DATEADD|DATEDIFF|CAST.*GETDATE', sql_upper):
|
||
tags.append('date_function')
|
||
|
||
return list(set(tags))
|
||
|
||
def _infer_difficulty(self, sql: str, tags: List[str]) -> str:
|
||
"""推断难度"""
|
||
sql_upper = sql.upper()
|
||
|
||
# 简单规则
|
||
if 'UNION' in sql_upper or 'EXCEPT' in sql_upper or 'INTERSECT' in sql_upper:
|
||
return 'hard'
|
||
if sql.count('JOIN') >= 3:
|
||
return 'hard'
|
||
if 'GROUP BY' in sql_upper and 'HAVING' in sql_upper:
|
||
return 'hard'
|
||
if re.search(r'SUBQUERY|EXISTS.*SELECT', sql_upper, re.IGNORECASE):
|
||
return 'hard'
|
||
if len(tags) >= 5:
|
||
return 'hard'
|
||
if sql.count('JOIN') == 2 or 'GROUP BY' in sql_upper:
|
||
return 'medium'
|
||
return 'easy'
|
||
|
||
def _extract_schema_from_sql(self, sql: str) -> Dict[str, Any]:
|
||
"""从SQL中提取表结构信息"""
|
||
import sqlglot
|
||
from sqlglot import exp
|
||
|
||
schema = {
|
||
"tables": {},
|
||
"joins": []
|
||
}
|
||
|
||
try:
|
||
parsed = sqlglot.parse_one(sql, dialect='tsql')
|
||
|
||
# 提取所有表
|
||
for table in parsed.find_all(exp.Table):
|
||
table_name = table.name
|
||
alias = table.alias_or_name
|
||
schema["tables"][table_name] = {
|
||
"alias": alias,
|
||
"columns": []
|
||
}
|
||
|
||
# 提取所有列引用
|
||
for col in parsed.find_all(exp.Column):
|
||
table_name = col.table
|
||
col_name = col.name
|
||
if table_name in schema["tables"]:
|
||
if col_name not in schema["tables"][table_name]["columns"]:
|
||
schema["tables"][table_name]["columns"].append(col_name)
|
||
|
||
# 提取JOIN关系
|
||
for join in parsed.find_all(exp.Join):
|
||
schema["joins"].append({
|
||
"type": join.args.get('side', 'INNER'),
|
||
"on": str(join.args.get('on'))
|
||
})
|
||
|
||
except Exception as e:
|
||
print(f"SQL解析失败: {e}")
|
||
|
||
return schema
|
||
|
||
def save_jsonl(self, output_path: Path):
|
||
"""保存为JSONL格式"""
|
||
with output_path.open('w', encoding='utf-8') as f:
|
||
for sample in self.samples:
|
||
f.write(json.dumps(sample.to_dict(), ensure_ascii=False) + '\n')
|
||
print(f"[OK] 保存 {len(self.samples)} 条样本到 {output_path}")
|
||
|
||
def save_fewshot_prompt(self, output_path: Path, limit_per_tag: int = 2):
|
||
"""生成few-shot提示文件"""
|
||
examples_by_tag = {}
|
||
for sample in self.samples:
|
||
for tag in sample.tags:
|
||
if tag not in examples_by_tag:
|
||
examples_by_tag[tag] = []
|
||
if len(examples_by_tag[tag]) < limit_per_tag:
|
||
examples_by_tag[tag].append(sample)
|
||
|
||
with output_path.open('w', encoding='utf-8') as f:
|
||
f.write("# Few-shot Examples for Text2SQL\n\n")
|
||
f.write("以下是从经验数据集中精选的示例，用于辅助SQL生成。\n\n")
|
||
|
||
for tag, samples in sorted(examples_by_tag.items()):
|
||
f.write(f"## 标签: {tag}\n\n")
|
||
for i, sample in enumerate(samples, 1):
|
||
f.write(f"### 示例 {i}\n")
|
||
f.write(f"**问题**：{sample.question_zh}\n\n")
|
||
f.write(f"**SQL**：\n```sql\n{sample.sql}\n```\n\n")
|
||
f.write(f"**说明**：{sample.explanation[:200]}...\n\n")
|
||
f.write("\n")
|
||
|
||
print(f"[OK] Few-shot示例已保存到 {output_path}")
|
||
|
||
def save_by_category(self, output_dir: Path):
|
||
"""按标签分类保存"""
|
||
output_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
by_tag = {}
|
||
for sample in self.samples:
|
||
for tag in sample.tags:
|
||
if tag not in by_tag:
|
||
by_tag[tag] = []
|
||
by_tag[tag].append(sample)
|
||
|
||
for tag, samples in by_tag.items():
|
||
tag_file = output_dir / f"{tag}.jsonl"
|
||
with tag_file.open('w', encoding='utf-8') as f:
|
||
for sample in samples:
|
||
f.write(json.dumps(sample.to_dict(), ensure_ascii=False) + '\n')
|
||
|
||
print(f"[OK] 按标签分类保存到 {output_dir}（{len(by_tag)}个类别）")
|
||
|
||
def generate_stats(self) -> Dict:
|
||
"""生成数据集统计"""
|
||
stats = {
|
||
"total_samples": len(self.samples),
|
||
"by_difficulty": {},
|
||
"by_tag": {},
|
||
"by_rating": {},
|
||
"avg_rating": 0,
|
||
}
|
||
|
||
# 难度分布
|
||
for diff in ['easy', 'medium', 'hard']:
|
||
stats["by_difficulty"][diff] = sum(
|
||
1 for s in self.samples if s.difficulty == diff
|
||
)
|
||
|
||
# 标签分布
|
||
tag_counts = {}
|
||
for sample in self.samples:
|
||
for tag in sample.tags:
|
||
tag_counts[tag] = tag_counts.get(tag, 0) + 1
|
||
stats["by_tag"] = tag_counts
|
||
|
||
# 评分分布
|
||
ratings = [s.rating for s in self.samples if s.rating is not None]
|
||
if ratings:
|
||
stats["avg_rating"] = sum(ratings) / len(ratings)
|
||
for r in range(1, 11):
|
||
stats["by_rating"][r] = ratings.count(r)
|
||
|
||
return stats
|
||
|
||
|
||
def main():
|
||
md_path = Path(__file__).parent.parent / "data" / "Example" / "Example_text2sql.md"
|
||
output_dir = Path(__file__).parent.parent / "data" / "experiences"
|
||
|
||
print("=" * 60)
|
||
print("Text2SQL 经验数据集构建工具")
|
||
print("=" * 60)
|
||
|
||
parser = ExampleText2SQLParser(md_path)
|
||
samples = parser.parse()
|
||
|
||
print(f"\n[OK] 解析完成，共 {len(samples)} 个样本")
|
||
|
||
# 生成统计
|
||
stats = parser.generate_stats()
|
||
print(f"\n[INFO] 数据集统计:")
|
||
print(f" 总样本数: {stats['total_samples']}")
|
||
print(f" 难度分布: {stats['by_difficulty']}")
|
||
print(f" 平均评分: {stats['avg_rating']:.1f}/10")
|
||
print(f"\n 标签分布（Top 10）:")
|
||
sorted_tags = sorted(stats["by_tag"].items(), key=lambda x: -x[1])[:10]
|
||
for tag, count in sorted_tags:
|
||
print(f" {tag}: {count}")
|
||
|
||
# 保存输出
|
||
output_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
# 1. JSONL格式（所有样本）
|
||
parser.save_jsonl(output_dir / "all_samples.jsonl")
|
||
|
||
# 2. 按难度分类
|
||
for diff in ['easy', 'medium', 'hard']:
|
||
diff_samples = [s for s in samples if s.difficulty == diff]
|
||
if diff_samples:
|
||
with (output_dir / f"difficulty_{diff}.jsonl").open('w', encoding='utf-8') as f:
|
||
for s in diff_samples:
|
||
f.write(json.dumps(s.to_dict(), ensure_ascii=False) + '\n')
|
||
|
||
# 3. Few-shot提示文件
|
||
parser.save_fewshot_prompt(output_dir / "fewshot_examples.md")
|
||
|
||
# 4. 按标签分类
|
||
parser.save_by_category(output_dir / "by_tag")
|
||
|
||
# 5. 高评分样本（rating >= 8）
|
||
high_rating = [s for s in samples if s.rating and s.rating >= 8]
|
||
with (output_dir / "high_rating_samples.jsonl").open('w', encoding='utf-8') as f:
|
||
for s in high_rating:
|
||
f.write(json.dumps(s.to_dict(), ensure_ascii=False) + '\n')
|
||
print(f"\n[OK] 高评分样本（rating>=8）: {len(high_rating)} 条")
|
||
|
||
# 6. 生成README
|
||
with (output_dir / "README.md").open('w', encoding='utf-8') as f:
|
||
f.write("# Text2SQL 经验数据集\n\n")
|
||
f.write(f"**来源**: `Example_text2sql.md`\n")
|
||
f.write(f"**样本数**: {len(samples)}\n")
|
||
f.write(f"**生成时间**: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n")
|
||
f.write("## 文件说明\n\n")
|
||
f.write("- `all_samples.jsonl` - 所有样本（JSONL格式）\n")
|
||
f.write("- `difficulty_*.jsonl` - 按难度分类\n")
|
||
f.write("- `by_tag/` - 按标签分类（每个标签一个文件）\n")
|
||
f.write("- `fewshot_examples.md` - Few-shot示例（Markdown格式）\n")
|
||
f.write("- `high_rating_samples.jsonl` - 高评分样本（rating≥8）\n\n")
|
||
f.write("## 使用建议\n\n")
|
||
f.write("1. **Few-shot prompting**: 使用 `fewshot_examples.md` 中的示例\n")
|
||
f.write("2. **训练数据**: 使用 `all_samples.jsonl` 进行模型微调\n")
|
||
f.write("3. **评估测试**: 使用 `high_rating_samples.jsonl` 作为测试集\n")
|
||
f.write("4. **领域适应**: 按标签选择相关示例（如broker相关查询）\n")
|
||
|
||
print(f"\n[OK] 所有文件已保存到: {output_dir}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|