wehub-resource-sync
296 lines
10 KiB
Python
296 lines
10 KiB
Python
import pandas as pd
|
|
import re
|
|
import os
|
|
import sys
|
|
from tqdm import tqdm
|
|
import requests
|
|
import json
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
|
|
# 添加项目根目录到路径,以便导入 config
|
|
script_dir = os.path.dirname(os.path.abspath(__file__))
|
|
project_root = os.path.dirname(os.path.dirname(script_dir))
|
|
if project_root not in sys.path:
|
|
sys.path.insert(0, project_root)
|
|
|
|
from config import settings
|
|
|
|
|
|
# 正则表达式模式,用于解析 QA 对
|
|
pattern = re.compile(r"-?\s*question: (.*?)\s*-?\s*answer: (.+)", re.DOTALL)
|
|
|
|
|
|
def generate(prompt, model_name=None, system=None, temperature=None,
|
|
max_tokens=None, stream=False, callback=None, api_key=None, base_url=None):
|
|
"""
|
|
使用 OpenAI API 格式生成响应
|
|
|
|
参数:
|
|
prompt: 用户输入的提示词
|
|
model_name: 模型名称(可选,默认从环境变量读取)
|
|
system: 系统提示词(可选)
|
|
temperature: 温度参数(可选,0-2之间)
|
|
max_tokens: 最大token数(可选)
|
|
stream: 是否流式输出(默认False)
|
|
callback: 回调函数(可选)
|
|
api_key: API密钥(可选,默认从配置读取)
|
|
base_url: API基础URL(可选,默认从配置读取)
|
|
|
|
返回:
|
|
full_response: 完整的响应文本
|
|
usage_info: token使用信息
|
|
"""
|
|
try:
|
|
# 从配置读取或使用参数
|
|
api_key = api_key or settings.OPENAI_API_KEY
|
|
base_url = base_url or settings.OPENAI_BASE_URL
|
|
default_model = "qwen/qwen3-vl-235b-a22b-instruct"
|
|
|
|
url = f"{base_url}/chat/completions"
|
|
|
|
# 构建消息列表
|
|
messages = []
|
|
if system:
|
|
messages.append({"role": "system", "content": system})
|
|
messages.append({"role": "user", "content": prompt})
|
|
|
|
# 构建请求payload
|
|
payload = {
|
|
"model": model_name or default_model,
|
|
"messages": messages,
|
|
"stream": stream
|
|
}
|
|
|
|
# 添加可选参数
|
|
if temperature is not None:
|
|
payload["temperature"] = temperature
|
|
if max_tokens is not None:
|
|
payload["max_tokens"] = max_tokens
|
|
|
|
# 设置请求头
|
|
headers = {
|
|
"Authorization": f"Bearer {api_key}",
|
|
"Content-Type": "application/json"
|
|
}
|
|
|
|
with requests.post(url, json=payload, headers=headers, stream=stream) as response:
|
|
response.raise_for_status()
|
|
|
|
full_response = ""
|
|
usage_info = None
|
|
|
|
if stream:
|
|
# 流式响应处理
|
|
for line in response.iter_lines():
|
|
if line:
|
|
line = line.decode('utf-8')
|
|
|
|
# 跳过注释行
|
|
if line.startswith(':'):
|
|
continue
|
|
|
|
# 移除 "data: " 前缀
|
|
if line.startswith('data: '):
|
|
line = line[6:]
|
|
|
|
# 检查是否结束
|
|
if line == '[DONE]':
|
|
break
|
|
|
|
try:
|
|
chunk = json.loads(line)
|
|
|
|
# 如果提供了回调函数,调用它
|
|
if callback:
|
|
callback(chunk)
|
|
else:
|
|
# 提取内容
|
|
if 'choices' in chunk and len(chunk['choices']) > 0:
|
|
delta = chunk['choices'][0].get('delta', {})
|
|
content = delta.get('content', '')
|
|
if content:
|
|
full_response += content
|
|
print(content, end="", flush=True)
|
|
|
|
# 获取使用信息(通常在最后一个chunk)
|
|
if 'usage' in chunk:
|
|
usage_info = chunk['usage']
|
|
|
|
except json.JSONDecodeError:
|
|
continue
|
|
|
|
if not callback:
|
|
print() # 换行
|
|
else:
|
|
# 非流式响应处理
|
|
result = response.json()
|
|
|
|
if callback:
|
|
callback(result)
|
|
else:
|
|
if 'choices' in result and len(result['choices']) > 0:
|
|
full_response = result['choices'][0]['message']['content']
|
|
print(full_response)
|
|
|
|
usage_info = result.get('usage')
|
|
|
|
return full_response, usage_info
|
|
|
|
except requests.exceptions.RequestException as e:
|
|
print(f"请求错误: {e}")
|
|
return None, None
|
|
except Exception as e:
|
|
print(f"发生错误: {e}")
|
|
return None, None
|
|
|
|
|
|
def process_single_row(index, row, file_name, rounds=3):
|
|
"""
|
|
处理单行数据，生成QA对
|
|
|
|
参数:
|
|
index: 行索引
|
|
row: DataFrame的行数据
|
|
file_name: 文件名（用于日志）
|
|
rounds: 生成轮次(默认3轮,选择最长结果)
|
|
|
|
返回:
|
|
tuple: (index, row_data) - 索引和包含Question/Answer的行数据字典
|
|
"""
|
|
from prompts import SYS_ED_TEMPLATE, ED_TEMPLATE
|
|
|
|
text = row['Text_pure']
|
|
if not isinstance(text, str):
|
|
print(f"Skipping row {index} in file {file_name}: Text is not a string")
|
|
row_data = row.to_dict()
|
|
row_data['Question'] = ''
|
|
row_data['Answer'] = ''
|
|
return index, row_data
|
|
|
|
best_question = ''
|
|
best_answer = ''
|
|
best_length = 0
|
|
|
|
# 执行多轮生成，选择最长的结果
|
|
for _ in range(rounds):
|
|
sys_prompt = SYS_ED_TEMPLATE
|
|
prompt = ED_TEMPLATE.format(text=text)
|
|
|
|
response, _ = generate(
|
|
system=sys_prompt,
|
|
prompt=prompt,
|
|
temperature=0.7
|
|
)
|
|
|
|
if response:
|
|
match = pattern.search(response)
|
|
if match:
|
|
question = match.group(1).strip()
|
|
answer = match.group(2).strip()
|
|
length = len(question) + len(answer)
|
|
if length > best_length:
|
|
best_question = question
|
|
best_answer = answer
|
|
best_length = length
|
|
|
|
# 构建返回数据
|
|
row_data = row.to_dict()
|
|
row_data['Question'] = best_question
|
|
row_data['Answer'] = best_answer
|
|
|
|
return index, row_data
|
|
|
|
|
|
def process_csv_to_qa(input_folder, output_folder, rounds=3, max_workers=5):
|
|
"""
|
|
将 CSV 文件中的文本转换为 QA 对
|
|
|
|
参数:
|
|
input_folder: 输入文件夹路径
|
|
output_folder: 输出文件夹路径
|
|
rounds: 每行生成的轮次(默认3,选择最长结果)
|
|
max_workers: 最大并发线程数(默认5)
|
|
|
|
逻辑:
|
|
- 读取所有 .csv 文件
|
|
- 对每行的 Text_pure 列生成 Question 和 Answer
|
|
- 使用多线程并发处理
|
|
- 支持断点续传(跳过已处理的行)
|
|
- 输出包含 Question 和 Answer 列的 CSV
|
|
"""
|
|
# 创建输出文件夹(如果不存在)
|
|
os.makedirs(output_folder, exist_ok=True)
|
|
|
|
# 获取所有CSV文件
|
|
csv_files = [f for f in os.listdir(input_folder) if f.endswith('.csv')]
|
|
|
|
for file_name in tqdm(csv_files, desc="Processing files"):
|
|
input_file_path = os.path.join(input_folder, file_name)
|
|
output_file_path = os.path.join(output_folder, file_name)
|
|
|
|
# 加载已有结果（如果存在）
|
|
if os.path.exists(output_file_path):
|
|
new_df = pd.read_csv(output_file_path, encoding='utf-8-sig')
|
|
else:
|
|
new_df = pd.DataFrame(columns=['Question', 'Answer'])
|
|
|
|
# 读取输入数据
|
|
df = pd.read_csv(input_file_path, encoding='utf-8-sig')
|
|
|
|
# 识别需要处理的行（跳过已完成的行）
|
|
rows_to_process = []
|
|
for index, row in df.iterrows():
|
|
# 检查是否已经处理过
|
|
if index < len(new_df) and pd.notnull(new_df.loc[index, 'Question']) and pd.notnull(new_df.loc[index, 'Answer']):
|
|
continue # 跳过已完成的行
|
|
rows_to_process.append((index, row))
|
|
|
|
if not rows_to_process:
|
|
print(f"File {file_name}: All rows already processed.")
|
|
continue
|
|
|
|
print(f"Processing {len(rows_to_process)} rows in {file_name} with {max_workers} workers...")
|
|
|
|
# 使用线程池并发处理
|
|
results = {} # 使用字典存储结果，key为index
|
|
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
|
# 提交所有任务
|
|
future_to_index = {
|
|
executor.submit(process_single_row, index, row, file_name, rounds): index
|
|
for index, row in rows_to_process
|
|
}
|
|
|
|
# 使用tqdm显示进度，并收集结果
|
|
for future in tqdm(as_completed(future_to_index), total=len(future_to_index), desc=f"Processing {file_name}"):
|
|
try:
|
|
index, row_data = future.result()
|
|
results[index] = row_data
|
|
except Exception as e:
|
|
index = future_to_index[future]
|
|
print(f"\nError processing row {index}: {e}")
|
|
# 发生错误时，创建空结果
|
|
results[index] = df.loc[index].to_dict()
|
|
results[index]['Question'] = ''
|
|
results[index]['Answer'] = ''
|
|
|
|
# 按索引顺序更新DataFrame
|
|
for index in sorted(results.keys()):
|
|
row_data = results[index]
|
|
if index < len(new_df):
|
|
new_df.iloc[index] = row_data
|
|
else:
|
|
new_df = pd.concat([new_df, pd.DataFrame([row_data])], ignore_index=True)
|
|
|
|
# 保存最终结果
|
|
new_df.to_csv(output_file_path, index=False, encoding='utf-8-sig')
|
|
print(f"File {file_name} processed and saved to {output_file_path}.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# 定义文件夹路径
|
|
input_folder = 'output'
|
|
output_folder = 'QA'
|
|
|
|
# 执行处理
|
|
process_csv_to_qa(input_folder, output_folder, rounds=3, max_workers=5)
|