巧用腾讯云OCR--图片组成的PDF也能构建知识库 | 技术创作特训营第一期

import json, PyPDF2, io, base64
from pdf2image import convert_from_path
from tencentcloud.common import credential
from tencentcloud.common.profile.client_profile import ClientProfile
from tencentcloud.common.profile.http_profile import HttpProfile
from tencentcloud.common.exception.tencent_cloud_sdk_exception import TencentCloudSDKException
from tencentcloud.ocr.v20181119 import ocr_client, models
from langchain import PromptTemplate
from langchain.chat_models import ChatOpenAI
from langchain.chains import LLMChain

def pdf_to_text(pdf_file):
    # 读取PDF文件
    pdf_reader = PyPDF2.PdfReader(pdf_file)
    print(pdf_reader.pages)
    num_pages = len(pdf_reader.pages)

    # 实例化一个认证对象
    SecretId = "XXXX"
    SecretKey = " XXXX"
    cred = credential.Credential(SecretId, SecretKey)
    httpProfile = HttpProfile()
    httpProfile.endpoint = "ocr.tencentcloudapi.com"
    clientProfile = ClientProfile()
    clientProfile.httpProfile = httpProfile
    client = ocr_client.OcrClient(cred, "ap-shanghai", clientProfile)

    # 处理每一页
    result_text = []
    for page_num in range(num_pages):
        images = convert_from_path(pdf_file, first_page=page_num+1, last_page=page_num+1)
        if images:
            image = images[0]
            buffer = io.BytesIO()
            image.save(buffer, format='JPEG')
            img_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')

            req = models.GeneralAccurateOCRRequest()
            params = {
                "ImageBase64": img_base64
            }
            req.from_json_string(json.dumps(params))

            try:
                resp = client.GeneralAccurateOCR(req)
                text = "\n".join([item.DetectedText for item in resp.TextDetections])
                result_text.append(text)
            except TencentCloudSDKException as err:
                print(f"Error in page {page_num}: {err}")

    return "\n".join(result_text)

def req_chat(question, prompt):
    llm = ChatOpenAI(openai_api_key=api_key, temperature=0.5, model_name='gpt-3.5-turbo')
    chain = LLMChain(llm = llm, prompt = prompt)
    
    return chain.run(question)

def gen_ret(pdf_text, prompt_template):
    prompt = PromptTemplate(
        input_variables=["text"],
        template=prompt_template,
    )

    rep = req_chat(pdf_text, prompt)

    return rep

summarize_template = """
    请总结下面三个反引号引起来的内容
    ```{text}```
    """
pdf_text = pdf_to_text(pdf_file)
print(gen_ret(pdf_text, summarize_template))

order_template = """
    从下面三个反引号引起来的内容，提取订并输出一个表格
    ```{text}```
    """
pdf_text = pdf_to_text(pdf_file)
print(gen_ret(pdf_text, order_template))