使用chatglm API处理论文
在这里chatglm链接注册并得到自己的密钥,一般会送个几百万token
背景:
由于毕设原因,得到了很多要读的论文,但是很多的pdf名称和论文实际名称不符,一个一个点开看有点太费力,刚好之前的项目有使用过chatglm-4的接口阅读文献,这里就修改下拿来读论文。
思路如下:
- 得到一个文件夹中所有的pdf文件目录
- 依次传给chatglm
- 生成下面的内容到excel
- 文献名称
- 作者
- 年份
- 类型
- 研究内容
- 方法
- 结果
- 讨论
1.整理成目录
得到一个文件夹中所有的pdf文件目录,目的是不会重复上传pdf
#1
import os
def get_all_pdf_filenames(directory):
# 初始化一个空列表,用于存储找到的 PDF 文件名
pdf_filenames = []
# 遍历给定目录及其所有子目录
for root, dirs, files in os.walk(directory):
# 为当前目录创建一个输出文件路径,文件名为 'pdf_filenames.txt'
output_file = os.path.join(root, 'pdf_filenames.txt')
# 遍历当前目录下的所有文件
for file in files:
# 检查文件是否以 '.pdf' 结尾
if file.endswith('.pdf'):
# 如果是 PDF 文件,将其文件名添加到列表中
pdf_filenames.append(file)
# 检查是否找到了 PDF 文件
if pdf_filenames:
# 如果列表不为空,则调用函数将文件名写入到指定的文本文件中
write_filenames_to_txt(pdf_filenames, output_file)
# 打印出写入文件的数量和路径
print(f'已将 {len(pdf_filenames)} 个 PDF 文件名写入 {output_file}')
# 清空列表以便于下一个目录使用
pdf_filenames.clear()
def write_filenames_to_txt(pdf_filenames, output_file):
# 使用 'with' 语句打开指定的输出文件进行写入,确保文件以 UTF-8 编码
with open(output_file, 'w', encoding='utf-8') as f:
# 遍历 PDF 文件名列表
for filename in pdf_filenames:
# 将每个文件名写入文件,并在每个文件名后添加换行符
f.write(filename + '\n')
def main():
# 提示用户输入要搜索的文件夹路径
directory = input("请输入要搜索的文件夹路径: ")
# 打印出用户输入的文件夹路径
print(directory)
# 调用函数以搜索文件夹并处理 PDF 文件
get_all_pdf_filenames(directory)
if __name__ == '__main__':
main()
2.提问并生成excel
查看chatglm4对应的文档,意识到不能直接上传pdf文件,需要先提取PDF文件中的文本,然后再结合prompt上传
格式大概如下:
from zhipuai import ZhipuAI
client = ZhipuAI(api_key="") # 填写您自己的APIKey
response = client.chat.completions.create(
model="glm-4-plus", # 填写需要调用的模型编码
messages=[
{"role": "system", "content": "你是一个乐于解答各种问题的助手,你的任务是为用户提供专业、准确、有见地的建议。"},
{"role": "user", "content": "农夫需要把狼、羊和白菜都带过河,但每次只能带一样物品,而且狼和羊不能单独相处,羊和白菜也不能单独相处,问农夫该如何过河。"}
],
)
print(response.choices[0].message)
首先循环读取pdf列表,得到pdf对应的path和当前的name
#读取文章列表
def generate_path(path):
#全局变量用于最后存入到excel
global df
#第一步生成的pdf文件名txt
work_station = path + '\\' + 'pdf_filenames.txt'
with open(work_station, 'r+', encoding='UTF-8') as f:
tasks = f.readlines()
while len(tasks) > 0:
filename=tasks[0].rstrip()
task = path + '\\' + filename
#如果有嵌套文件夹
if os.path.isdir(task):
generate_path(task)
#没有就处理单个文件
else:
do_task(task,filename)
print(task)
print(filename)
tasks = tasks[1:]
f.truncate(0)
#处理完全部的就生成到excel中
df.to_excel(path + '\\' +"literature_info.xlsx", index=False)
接下来就是处理单个pdf,提取PDF文件中的文本
# 提取PDF文件中的文本
# 记得pip install PdfReader
def extract_text_from_pdf(pdf_file_path):
reader = PdfReader(pdf_file_path)
text = ""
for page in reader.pages:
text += page.extract_text()
return text
然后按照文档给的模板 调用chatGLM-4 API
def call_chat_glm_api(text):
content=gen_prompt(text)
try:
response = client.chat.completions.create(
model="glm-4",
messages=[
{"role": "user", "content": content}
],
)
except Exception as e:
print(f"发生错误:{e}")
return none
return response.choices[0].message
提示词就是仿照之前的模板,注意要求使用中文回答
def gen_prompt(text: str) -> str:
prompt = "你是一名密码分析领域的专家。现在有一篇关于密码分析的文章,内容为:\n[" + text + "]\n"
prompt += ("请根据上述文章的内容,用中文回答,依次生成 文献名称、作者、年份、类型、研究内容、方法、结果、讨论\n")
prompt+=("使用中文回答,回答考虑下面三点:\n"+
"1.不要使用1...2...3...4..回答\n"+
"2.给出一段话,不要给出一句简短的回答。\n"+
"3.研究内容、方法、结果和讨论的回答尽可能详细充实,每个都要大于70个字符。\n")
prompt += "请按照以下格式生成中文的回答:\n"
prompt += "文献名称:......\n作者:......\n年份:......\n类型:......\n研究内容:......\n方法:......\n结果:......\n讨论:......\n"
return prompt
接下来就是对单个pdf进行提问
def do_task(Filepath,filename):
pdf_file_path = Filepath
#提取文本
extracted_text = extract_text_from_pdf(pdf_file_path)
#得到回答
api_response = call_chat_glm_api(extracted_text)
#将得到的回答转为text格式
text = api_response.content if hasattr(api_response, 'content') and isinstance(api_response.content, str) else ""
print(filename +" 已完成\n "+text)
#添加入excel
add2Excel(text,filename)
解析text并将其添加到excel中
def add2Excel(text,filename):
# 有时候返回的信息会是全英文 所以需要none作特别处理
global df
# 提取信息
info = {
"pdf名称": filename,
"文献名称": text.split("文献名称")[1].split("\n")[0] if "文献名称" in text else None,
"作者": text.split("作者")[1].split("\n")[0] if "作者" in text else None,
"年份": text.split("年份")[1].split("\n")[0] if "年份" in text else None,
"类型": text.split("类型")[1].split("\n")[0] if "类型" in text else None,
"研究内容": text.split("研究内容")[1].split("\n")[0] if "研究内容" in text else None,
"方法": text.split("方法")[1].split("\n")[0] if "方法" in text else None,
"结果": text.split("结果")[1].split("\n")[0] if "结果" in text else None,
"讨论": text.split("讨论")[1].split("\n")[0] if "讨论" in text else None,
}
# 尝试读取现有的Excel文件,如果不存在则创建一个新的DataFrame
# df = df.append(info, ignore_index=True)
# 将新信息添加到DataFrame
df = pd.concat([df, pd.DataFrame([info])], ignore_index=True)
最后就是main
# 保存原始stdout引用
original_stdout = sys.stdout
# 主流程
if __name__ == '__main__':
global root
args = sys.argv
root = args[1]
key = args[2]
client = ZhipuAI(api_key=key)
# 打开一个文件用于写入print方便保存log
with open(root+ '\\' +'output.txt', 'w', encoding='utf-8') as f:
# 将stdout重定向到文件
sys.stdout = f
# 调用函数,其print输出将被重定向到文件
generate_path(root)
# 完成输出后,恢复原始stdout
sys.stdout = original_stdout
3.运行
实际运行中发现上面的add2Excel
函数可能无法完整的提取回答,比如回答为
结果: 对于LEA,作者找到了更高效的12轮和13轮差分特征,其概率比之前最好的12轮和13轮差分特征好约26倍和27倍。对于HIGHT,他们找到了新的12轮和13轮差分特征,尽管概率与之前报告的最佳概率相同。研究结果扩展了LEA和HIGHT密码的差分分析。
只能提取到
对于LEA,作者找到了更高效的12轮和13轮差分特征,其概率比之前最好的12轮和13轮差分特征好约26倍和27倍。对于HIGHT,他们找到了新的12轮和13轮差分特征,尽管概率与之前报告的最佳概率相同。研究
会缺失最后的部分,思考后改进如下,用下一个小标题去作为分行的区分,就不会出现识别的问题。
info = {
"pdf名称": filename,
"文献名称": text.split("文献名称:")[1].split("\n作者:")[0].strip() if "文献名称:" in text else None,
"作者": text.split("作者:")[1].split("\n年份:")[0].strip() if "作者:" in text else None,
"年份": text.split("年份:")[1].split("\n类型:")[0].strip() if "年份:" in text else None,
"类型": text.split("类型:")[1].split("\n研究内容:")[0].strip() if "类型:" in text else None,
"研究内容": text.split("研究内容:")[1].split("\n方法:")[0].strip() if "研究内容:" in text else None,
"方法": text.split("方法:")[1].split("\n结果:")[0].strip() if "方法:" in text else None,
"结果": text.split("结果:")[1].split("\n讨论:")[0].strip() if "结果:" in text else None,
"讨论": text.split("讨论:")[1].strip() if "讨论:" in text else None,
}
还有一些文档在回答时只会返回英文,即使prompt里明确指出使用中文回答也没有用,这里选择跳过。
使用方法:
首先使用1生成对应目录的pdf_filenames.txt
,然后使用2生成对应目录的literature_info.xlsx
即可,对于excel中没有得到回答的paper,可以在output中查看对应的英文回答。
还有一个读取单个文献的小demo,但是这样的话其实使用UI界面会更快,就没必要放出来了