【Datawhale AI夏令营】科大讯飞AI大赛（大模型技术）/夏令营：让AI理解列车排期表（Task3）

我没招了jpg

import pandas as pd
import requests
import re
import json
from tqdm import tqdm
from datetime import datetime, timedeltadef calculate_stop_duration(arrival_time_str, departure_time_str):"""计算列车停留时长，处理跨天和异常情况"""try:if arrival_time_str == '无数据' or departure_time_str == '无数据':return "无数据"# 统一时间格式（处理"23:41"和"23:41:00"混用情况）arrival_time_str = arrival_time_str if len(arrival_time_str.split(':')) == 3 else arrival_time_str + ':00'departure_time_str = departure_time_str if len(departure_time_str.split(':')) == 3 else departure_time_str + ':00'arrival_time = datetime.strptime(arrival_time_str, '%H:%M:%S')departure_time = datetime.strptime(departure_time_str, '%H:%M:%S')# 处理跨天情况（如23:40到达，00:12离开）if departure_time < arrival_time:departure_time += timedelta(days=1)duration = departure_time - arrival_timehours, remainder = divmod(duration.seconds, 3600)minutes = remainder // 60# 优化输出格式（小于1小时只显示分钟）if hours == 0:return f"{minutes}分钟"else:return f"{hours}小时{minutes}分钟"except Exception as e:return f"计算错误: {str(e)}"def call_llm(content: str):"""调用大模型API（与原baseline相同）"""url = "https://api.siliconflow.cn/v1/chat/completions"payload = {"model": "Qwen/Qwen3-8B","messages": [{"role": "user", "content": content}]}headers = {"Authorization": "Bearer sk-xxx","Content-Type": "application/json"}resp = requests.request("POST", url, json=payload, headers=headers).json()content = resp['choices'][0]['message']['content'].split('</think>')[-1]pattern = re.compile(r'^```json\s*([\s\S]*?)```$', re.IGNORECASE)match = pattern.match(content.strip())if match:return json.loads(match.group(1).strip())return contentdef create_question_list(row: dict):"""生成多样化问题列表（包含简单查询和时间推理）"""question_list = []train_number = row["车次"]departure_station = row["始发站"]# 1. 基础信息查询question_list.append(f'{train_number}号车次应该从哪个检票口检票？')question_list.append(f'{train_number}号车次应该从哪个站台上车？')question_list.append(f'{train_number}次列车的终到站是哪里？')question_list.append(f'{train_number}次列车的候车区域是哪里？')# 2. 时间推理问题if row['到点'] != '无数据' and row['开点'] != '无数据':# 停留时长question_list.append(f'{train_number}次列车在{departure_station}站的停留时长是多久？')# 发车时间查询（验证模型能否处理时间格式）question_list.append(f'{train_number}次列车什么时候从{departure_station}站发车？')# 3. 比较类问题（需要多行数据，可在后续扩展）# question_list.append(f'从{departure_station}出发的车次中，哪趟停留时间最长？')return question_list# 主流程
if __name__ == "__main__":# 1. 读取数据data = pd.read_excel('data/info_table.xlsx').fillna('无数据')# 2. 预处理时间数据data['停留时长'] = data.apply(lambda row: calculate_stop_duration(str(row['到点']), str(row['开点'])), axis=1)# 3. 生成提示词模板prompt = '''你是专业的列车乘务员，请严格根据以下列车信息回答问题：
# 列车信息
{}# 用户问题列表
{}注意：
1. 所有时间计算已考虑跨天情况
2. 对于"无数据"的问题，如实回答"暂无该信息"
'''output_format = '''# 输出格式
```json
[{"q": "问题文本","a": "答案文本（直接引用或计算得出）"
}]
```'''# 4. 生成训练数据train_data = []error_log = []for idx, row in tqdm(data.iterrows(), total=len(data)):try:row_dict = row.to_dict()questions = create_question_list(row_dict)# 调用LLM生成答案（使用原始数据+计算后的停留时长）llm_response = call_llm(prompt.format(row_dict, questions) + output_format)if isinstance(llm_response, list):train_data.extend([{"instruction": qa["q"], "output": qa["a"]} for qa in llm_response])else:error_log.append({"row": row_dict, "error": "LLM返回格式错误"})except Exception as e:error_log.append({"row": row_dict, "error": str(e)})# 5. 保存结果with open('train_data_v2.json', 'w', encoding='utf-8') as f:json.dump(train_data, f, ensure_ascii=False, indent=2)if error_log:with open('error_log.json', 'w', encoding='utf-8') as f:json.dump(error_log, f, ensure_ascii=False, indent=2)print(f"生成完成！有效数据{len(train_data)}条，错误{len(error_log)}条")