Add pdf、docx parser and image、video to text function for knowledge
This commit is contained in:
wugren
@@ -1,17 +1,30 @@
|
||||
import copy
|
||||
import os
|
||||
from typing import List
|
||||
|
||||
import aiofiles
|
||||
import chardet
|
||||
import logging
|
||||
import string
|
||||
from knowledge import ImageObjectBuilder, DocumentObjectBuilder, KnowledgePipelineEnvironment, KnowledgePipelineJournal
|
||||
from aios_kernel.storage import AIStorage
|
||||
|
||||
import docx2txt
|
||||
from PyPDF2 import PdfReader
|
||||
|
||||
from aios import KnowledgePipelineEnvironment, ImageObjectBuilder, DocumentObjectBuilder, KnowledgeStore, RichTextObject
|
||||
from aios.agent.agent_base import AgentPrompt
|
||||
from aios.frame.compute_kernel import ComputeKernel
|
||||
from aios.knowledge.data.writer import split_text
|
||||
from aios.proto.compute_task import ComputeTaskResult, ComputeTaskResultCode
|
||||
from aios.storage.storage import AIStorage
|
||||
from aios.utils import video_utils, image_utils
|
||||
|
||||
|
||||
class KnowledgeDirSource:
|
||||
def __init__(self, env: KnowledgePipelineEnvironment, config):
|
||||
self.env = env
|
||||
path = string.Template(config["path"]).substitute(myai_dir=AIStorage.get_instance().get_myai_dir())
|
||||
config["path"] = path
|
||||
self.config = config
|
||||
self.config = config
|
||||
|
||||
# @classmethod
|
||||
# def user_config_items(cls):
|
||||
@@ -19,16 +32,16 @@ class KnowledgeDirSource:
|
||||
|
||||
def path(self):
|
||||
return self.config["path"]
|
||||
|
||||
|
||||
@staticmethod
|
||||
async def read_txt_file(file_path:str)->str:
|
||||
cur_encode = "utf-8"
|
||||
async with aiofiles.open(file_path,'rb') as f:
|
||||
cur_encode = chardet.detect(await f.read())['encoding']
|
||||
|
||||
|
||||
async with aiofiles.open(file_path,'r',encoding=cur_encode) as f:
|
||||
return await f.read()
|
||||
|
||||
|
||||
async def next(self):
|
||||
while True:
|
||||
journals = self.env.journal.latest_journals(1)
|
||||
@@ -42,7 +55,7 @@ class KnowledgeDirSource:
|
||||
if os.path.getmtime(self.path()) <= from_time:
|
||||
yield (None, None)
|
||||
continue
|
||||
|
||||
|
||||
file_pathes = sorted(os.listdir(self.path()), key=lambda x: os.path.getctime(os.path.join(self.path(), x)))
|
||||
for rel_path in file_pathes:
|
||||
file_path = os.path.join(self.path(), rel_path)
|
||||
@@ -62,7 +75,138 @@ class KnowledgeDirSource:
|
||||
await self.env.get_knowledge_store().insert_object(document)
|
||||
yield (document.calculate_id(), file_path)
|
||||
yield (None, None)
|
||||
|
||||
|
||||
|
||||
def init(env: KnowledgePipelineEnvironment, params: dict) -> KnowledgeDirSource:
|
||||
return KnowledgeDirSource(env, params)
|
||||
return KnowledgeDirSource(env, params)
|
||||
|
||||
|
||||
async def image_to_text(images: List[str]) -> str:
|
||||
msg_prompt = AgentPrompt()
|
||||
image_prompt = "What's in this image?"
|
||||
content = [{"type": "text", "text": image_prompt}]
|
||||
content.extend([{"type": "image_url", "image_url": {"url": image_utils.to_base64(image)}} for image in images])
|
||||
msg_prompt.messages = [{"role": "user", "content": content}]
|
||||
|
||||
resp: ComputeTaskResult = await (ComputeKernel.get_instance()
|
||||
.do_llm_completion(prompt=msg_prompt,
|
||||
resp_mode="text",
|
||||
mode_name="gpt-4-vision-preview",
|
||||
max_token=4000,
|
||||
inner_functions=None,
|
||||
timeout=None))
|
||||
if resp.result_code != ComputeTaskResultCode.OK:
|
||||
raise Exception(f"image_to_text error: {resp.result_code} msg:{resp.error_str}")
|
||||
return resp.result_str
|
||||
|
||||
|
||||
async def video_to_text(video: str) -> str:
|
||||
prompt = "These pictures are key frames extracted from the video. Please describe the content of the video based on these key frames."
|
||||
frames = video_utils.extract_frames(video, (1024, 1024))
|
||||
msg_prompt = AgentPrompt()
|
||||
content = [{"type": "text", "text": prompt}]
|
||||
content.extend([{"type": "image_url", "image_url": {"url": frame}} for frame in frames])
|
||||
msg_prompt.messages = [{"role": "user", "content": content}]
|
||||
resp: ComputeTaskResult = await (ComputeKernel.get_instance()
|
||||
.do_llm_completion(prompt=msg_prompt,
|
||||
resp_mode="text",
|
||||
mode_name="gpt-4-vision-preview",
|
||||
max_token=4000,
|
||||
inner_functions=None,
|
||||
timeout=None))
|
||||
if resp.result_code != ComputeTaskResultCode.OK:
|
||||
raise Exception(f"video_to_text error: {resp.result_code} msg:{resp.error_str}")
|
||||
return resp.result_str
|
||||
|
||||
|
||||
async def summary_document(text: str, separators: List[str]=["\n\n", "\n"]) -> str:
|
||||
chunks = split_text(text, separators=separators, chunk_size=4000, chunk_overlap=200, length_function=len)
|
||||
|
||||
prompt = AgentPrompt()
|
||||
prompt.system_message = {"role":"system","content":"Your job is to generate a summary based on the input."}
|
||||
if len(chunks) == 1:
|
||||
prompt.append(AgentPrompt(chunks[0]))
|
||||
resp = await (ComputeKernel.get_instance()
|
||||
.do_llm_completion(prompt=prompt,
|
||||
resp_mode="text",
|
||||
mode_name="gpt-4-1106-preview",
|
||||
max_token=4000,
|
||||
inner_functions=None,
|
||||
timeout=None))
|
||||
if resp.result_code != ComputeTaskResultCode.OK:
|
||||
raise Exception(f"summary_document error: {resp.result_code} msg:{resp.error_str}")
|
||||
return resp.result_str
|
||||
|
||||
segments = []
|
||||
for i, chunk in enumerate(chunks):
|
||||
seg_prompt = copy.deepcopy(prompt)
|
||||
seg_prompt.append(AgentPrompt(chunk))
|
||||
resp = await (ComputeKernel.get_instance()
|
||||
.do_llm_completion(prompt=seg_prompt,
|
||||
resp_mode="text",
|
||||
mode_name="gpt-4-1106-preview",
|
||||
max_token=4000,
|
||||
inner_functions=None,
|
||||
timeout=None))
|
||||
if resp.result_code != ComputeTaskResultCode.OK:
|
||||
raise Exception(f"summary_document error: {resp.result_code} msg:{resp.error_str}")
|
||||
segments.append(resp.result_str)
|
||||
|
||||
segments_str = "\n".join(segments)
|
||||
prompt.append(AgentPrompt(f"Please combine the summaries of the following paragraphs into one complete summary:\n{segments_str}"))
|
||||
resp = await (ComputeKernel.get_instance()
|
||||
.do_llm_completion(prompt=prompt,
|
||||
resp_mode="text",
|
||||
mode_name="gpt-4-1106-preview",
|
||||
max_token=4000,
|
||||
inner_functions=None,
|
||||
timeout=None))
|
||||
if resp.result_code != ComputeTaskResultCode.OK:
|
||||
raise Exception(f"summary_document error: {resp.result_code} msg:{resp.error_str}")
|
||||
return resp.result_str
|
||||
|
||||
|
||||
|
||||
def pdf_to_rich_text_object(pdf: str, store: KnowledgeStore) -> RichTextObject:
|
||||
base_name = os.path.basename(pdf)
|
||||
cache_path = os.path.join(AIStorage.get_instance().get_myai_dir(), "knowledge", "doc_cache", base_name)
|
||||
if not os.path.exists(cache_path):
|
||||
os.makedirs(cache_path)
|
||||
|
||||
reader = PdfReader(pdf)
|
||||
rich_text = RichTextObject()
|
||||
page_texts = []
|
||||
image_count = 0
|
||||
for page in reader.pages:
|
||||
text = page.extract_text()
|
||||
page_texts.append(text)
|
||||
for image in page.images:
|
||||
image_path = os.path.join(cache_path, f"{image_count}_{image.name}")
|
||||
with open(image_path, "wb") as f:
|
||||
f.write(image.data)
|
||||
image_object = ImageObjectBuilder({}, {}, image_path).build(store)
|
||||
rich_text.add_image(image_object)
|
||||
|
||||
document = DocumentObjectBuilder({}, {}, "".join(page_texts)).build(store)
|
||||
rich_text.add_document(document)
|
||||
|
||||
return rich_text
|
||||
|
||||
|
||||
def doc_to_rich_text_object(doc: str, store: KnowledgeStore) -> RichTextObject:
|
||||
base_name = os.path.basename(doc)
|
||||
cache_path = os.path.join(AIStorage.get_instance().get_myai_dir(), "knowledge", "doc_cache", base_name)
|
||||
if not os.path.exists(cache_path):
|
||||
os.makedirs(cache_path)
|
||||
text = docx2txt.process(doc, cache_path)
|
||||
|
||||
rich_text = RichTextObject()
|
||||
for image in os.listdir(cache_path):
|
||||
image_path = os.path.join(cache_path, image)
|
||||
image_object = ImageObjectBuilder({}, {}, image_path).build(store)
|
||||
rich_text.add_image(image_object)
|
||||
|
||||
document = DocumentObjectBuilder({}, {}, text).build(store)
|
||||
rich_text.add_document(document)
|
||||
|
||||
return rich_text
|
||||
|
||||
Reference in New Issue
Block a user