158 lines
6.2 KiB
Python
158 lines
6.2 KiB
Python
from langchain_core.documents import Document
|
||
from typing import List
|
||
from langchain.document_loaders.unstructured import UnstructuredFileLoader
|
||
# 导入模块
|
||
import csv
|
||
import os
|
||
import re
|
||
import json
|
||
import pandas as pd
|
||
|
||
# TODO: 以CSV格式加载内容
|
||
MAX_CONTENT = 18000
|
||
SPECIAL_SYMBOL = ['\'', '"']
|
||
class RapidOCRCSVLoader(UnstructuredFileLoader):
|
||
|
||
def _get_elements(self) -> List:
|
||
meta_path = self.file_path.split("/")[-1]
|
||
|
||
# 读取csv文件
|
||
def read_csv_file(file_path, sep = ',', encoding = 'utf-8'):
|
||
df_data = pd.read_csv(file_path, sep = sep, encoding = encoding)
|
||
df_data = df_data.apply(remove_duplicate_pandas, axis = 1)
|
||
df_data = df_data.dropna(subset=['HTML全文'])
|
||
|
||
return df_data
|
||
|
||
def remove_duplicate_string(x):
|
||
if not isinstance(x, str):
|
||
return x
|
||
if len(x) <= 0:
|
||
return x
|
||
|
||
if x[0] in SPECIAL_SYMBOL:
|
||
x = x[1:]
|
||
|
||
if x[-1] in SPECIAL_SYMBOL:
|
||
x = x[:-1]
|
||
|
||
return x
|
||
|
||
def remove_duplicate_pandas(row):
|
||
headers = ['标题', 'HTML全文','发布时间', '发布机构来源', '发布机构',
|
||
'适用地区', '主题', '详情地址', '资源来源分类',
|
||
'资源来源名称', '资源来源标识']
|
||
|
||
for col in headers:
|
||
row[col] = remove_duplicate_string(row[col])
|
||
|
||
if len(row[col]) > MAX_CONTENT:
|
||
row[col] = row[col][:MAX_CONTENT]
|
||
|
||
return row
|
||
|
||
def apply_title_content(x):
|
||
# 处理 title相关的内容
|
||
titles = [x['标题']]
|
||
pattern = r'##title##(.*?)##/title##'
|
||
title_list = re.findall(pattern, x['HTML全文'])
|
||
titles += title_list
|
||
|
||
# 去除内容中标题的特殊符号
|
||
content = x['HTML全文'].replace('##title##', '').replace('##/title##', '').replace('\\u3000', '')
|
||
title_str = "\n".join(titles)
|
||
tilte_str = title_str.replace('\\u3000', '')
|
||
|
||
content = content.replace('#####', '\n')
|
||
|
||
# 把内容相关提出
|
||
title_idx = content.rfind("##/title##") + 1
|
||
|
||
return Document(page_content=x['标题'],
|
||
metadata={"path": meta_path,
|
||
"_type": 'title',
|
||
'title' : x['标题'],
|
||
"content": content ,
|
||
# 'file_number': x['发文文号'],
|
||
# 'keywords' : x['关键词'],
|
||
'release_date' : x['发布时间'],
|
||
'sourceOrganization' : x['发布机构来源'],
|
||
'organization' : x['发布机构'],
|
||
'region' : x['适用地区'],
|
||
'subject' : x['主题'],
|
||
'source' : x['详情地址'],
|
||
'datasource' : x['资源来源分类'],
|
||
'datasourceclass' : x['资源来源名称'],
|
||
'datasource_key' : x['资源来源标识'],
|
||
"xml": title_str,
|
||
"tag": tag})
|
||
|
||
# title作为query，content作为answer
|
||
def title_content2text(pd_obj, filepath, tag = 'csv'):
|
||
ret_list = []
|
||
pd_obj = pd_obj.apply(apply_title_content, axis=1)
|
||
ret_list = pd_obj.values.tolist()
|
||
return ret_list
|
||
|
||
def apply_content(x):
|
||
# 去除内容中标题的特殊符号
|
||
# 把内容相关提出
|
||
content = x['HTML全文']
|
||
title_idx = content.rfind("##/title##") + 1
|
||
if title_idx > 0 and title_idx < len(content):
|
||
content = content[title_idx + 1:]
|
||
content = content.replace('##title##', '').replace('##/title##', '').replace('\\u3000', '')
|
||
content = content.replace('#####', '\n')
|
||
|
||
# 对内容进行切分
|
||
return Document(page_content=content,
|
||
metadata={"path": meta_path,
|
||
"_type": 'content',
|
||
'title' : x['标题'],
|
||
"content": '' ,
|
||
# 'file_number': x['发文文号'],
|
||
# 'keywords' : x['关键词'],
|
||
'release_date' : x['发布时间'],
|
||
'sourceOrganization' : x['发布机构来源'],
|
||
'organization' : x['发布机构'],
|
||
'region' : x['适用地区'],
|
||
'subject' : x['主题'],
|
||
'source' : x['详情地址'],
|
||
'datasource' : x['资源来源分类'],
|
||
'datasourceclass' : x['资源来源名称'],
|
||
'datasource_key' : x['资源来源标识'],
|
||
"xml": '',
|
||
"tag": tag})
|
||
|
||
def content2text(pd_obj, filepath, tag = 'csv'):
|
||
ret_list = []
|
||
pd_obj = pd_obj.apply(apply_content, axis=1)
|
||
ret_list = pd_obj.values.tolist()
|
||
return ret_list
|
||
|
||
# 读取 csv 文件
|
||
pd_obj = read_csv_file(self.file_path, r'\^\|\|\^')
|
||
tag = "csv"
|
||
|
||
# 解析文件，获取 query(标题)-value(全文)
|
||
ret_list_title = title_content2text(pd_obj, self.file_path, 'csv')
|
||
# 解析文件，获取内容，并进行切割
|
||
ret_list_content = content2text(pd_obj, self.file_path, 'csv')
|
||
ret_list = ret_list_title + ret_list_content
|
||
|
||
return ret_list
|
||
|
||
def load(self) -> List[Document]:
|
||
return self._get_elements()
|
||
|
||
if __name__ == '__main__':
|
||
loader = RapidOCRCSVLoader(file_path="//home/work/project/test_result.csv")
|
||
docs = loader.load()
|
||
|
||
i =0
|
||
for doc in docs:
|
||
|
||
print(doc)
|
||
i += 1
|
||
if i>1000:
|
||
break |