File size: 1,713 Bytes
c4ad2cf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 |
import pdfplumber
import pandas as pd
from openai.embeddings_utils import get_embedding
import openai
try:
from utils.config import OPENAI_KEY
except ImportError:
from utils.config_sample import OPENAI_KEY
class PdfFile:
'''
1. pdf文件解析器, 目前使用pdfplumber模块简单实现
2. 之后需要找更好的ML模块以能更好的提取文章内容(精确到节段), 位置定位更精确
3. 在有新的提取方式以前, 这部分不添加设置项
'''
def __init__(self, file, api_key=None):
self.pdf_file = pdfplumber.open(file)
self.api_key = api_key if api_key else OPENAI_KEY
def __del__(self):
if self.pdf_file:
self.pdf_file.close()
def parse_info(self):
'''
一页内容存一次, 所有内容通过模型进行预测
'''
all_text = []
for page in self.pdf_file.pages:
page_text = page.dedupe_chars().extract_text(layout=True)
if page_text:
all_text.append({
'page': page.page_number,
'text': page_text
})
text_data = pd.DataFrame(all_text)
openai.api_key = self.api_key
text_data['embedding'] = text_data['text'].apply(lambda x: get_embedding(x, engine='text-embedding-ada-002'))
# text_data.to_csv('debug.data.csv', index=0)
return text_data
@property
def metadata(self):
'''
调用rispy给出可解析的所有字段
'''
return self.pdf_file.metadata
if __name__ == '__main__':
pdf = PdfFile(file='/home/silen/git_proj/ReviewGPT/test/demo_paper.pdf')
print(pdf.parse_info())
|