|
|
|
|
|
from openai import OpenAI
|
|
|
|
|
|
|
|
|
|
|
|
from Config import MODEL_NAME, MODEL_API_KEY
|
|
|
|
|
|
|
|
|
|
|
|
# 初始化 OpenAI 客户端
|
|
|
|
|
|
client = OpenAI(
|
|
|
|
|
|
api_key=MODEL_API_KEY,
|
|
|
|
|
|
base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
|
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# 利用AI获取数据集的X轴和Y轴的列名
|
|
|
|
|
|
def generate_columns_with_ai(data):
|
|
|
|
|
|
"""
|
|
|
|
|
|
利用大模型解析数据，生成 category_columns_str 和 value_column_str
|
|
|
|
|
|
:param data: 数据集（列表字典格式，例如：[{"行政区划名": "二道区", "学校名称": "清华附中", "课程数量": 100}, ...]）
|
|
|
|
|
|
:param api_key: OpenAI API 密钥
|
|
|
|
|
|
:return: (category_columns_str, value_column_str)
|
|
|
|
|
|
"""
|
|
|
|
|
|
# 获取所有字段名
|
|
|
|
|
|
columns = list(data[0].keys()) if data else []
|
|
|
|
|
|
|
|
|
|
|
|
# 构造提示词
|
|
|
|
|
|
prompt = f"""
|
|
|
|
|
|
给定以下字段名列表：{columns}，请分析并回答以下问题：
|
|
|
|
|
|
1. 哪些字段适合作为分类字段（category_columns_str）？请用逗号分隔。
|
|
|
|
|
|
2. 哪个字段适合作为数值字段（value_column_str）？
|
|
|
|
|
|
3. 以JSON格式返回结果,但不要输出 json``` ，```
|
|
|
|
|
|
|
|
|
|
|
|
返回格式：
|
|
|
|
|
|
category_columns_str: <字段1,字段2,...>
|
|
|
|
|
|
value_column_str: <字段>
|
|
|
|
|
|
"""
|
|
|
|
|
|
|
|
|
|
|
|
# 调用大模型
|
|
|
|
|
|
response = client.chat.completions.create(
|
|
|
|
|
|
model=MODEL_NAME,
|
|
|
|
|
|
messages=[
|
|
|
|
|
|
{"role": "system",
|
|
|
|
|
|
"content": "你是一个专业的语义分类助手。"},
|
|
|
|
|
|
{"role": "user", "content": prompt}
|
|
|
|
|
|
],
|
|
|
|
|
|
max_tokens=500
|
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
# 解析模型返回的结果
|
|
|
|
|
|
result = response.choices[0].message.content
|
|
|
|
|
|
category_columns_str = result.split("category_columns_str: ")[1].split("\n")[0].strip()
|
|
|
|
|
|
value_column_str = result.split("value_column_str: ")[1].strip()
|
|
|
|
|
|
|
|
|
|
|
|
return category_columns_str, value_column_str
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# 示例数据
|
|
|
|
|
|
data = [
|
|
|
|
|
|
{"行政区划名": "二道区", "学校名称": "清华附中", "课程数量": 100},
|
|
|
|
|
|
{"行政区划名": "朝阳区", "学校名称": "复旦附中", "课程数量": 120},
|
|
|
|
|
|
]
|
|
|
|
|
|
|
|
|
|
|
|
# 调用函数
|
|
|
|
|
|
category_columns_str, value_column_str = generate_columns_with_ai(data)
|
|
|
|
|
|
print(f"category_columns_str: {category_columns_str}")
|
|
|
|
|
|
print(f"value_column_str: {value_column_str}")
|
