mirror of
https://github.com/HKUDS/RAG-Anything.git
synced 2026-07-07 22:54:45 +08:00
295 lines
8.8 KiB
Python
295 lines
8.8 KiB
Python
"""
|
|
Example of directly using modal processors
|
|
|
|
This example demonstrates how to use RAG-Anything's modal processors directly without going through MinerU.
|
|
Includes examples for:
|
|
- Image processing (with vision model)
|
|
- Table processing (with LLM)
|
|
- Equation processing (with LLM)
|
|
- Audio processing (with audio LLM)
|
|
"""
|
|
|
|
import asyncio
|
|
import argparse
|
|
from lightrag.llm.openai import openai_complete_if_cache, openai_embed
|
|
from lightrag.utils import EmbeddingFunc
|
|
from lightrag.kg.shared_storage import initialize_pipeline_status
|
|
from lightrag import LightRAG
|
|
from raganything.modalprocessors import (
|
|
ImageModalProcessor,
|
|
TableModalProcessor,
|
|
EquationModalProcessor,
|
|
AudioModalProcessor,
|
|
)
|
|
|
|
WORKING_DIR = "./rag_storage"
|
|
|
|
|
|
def get_llm_model_func(api_key: str, base_url: str = None):
|
|
return (
|
|
lambda prompt,
|
|
system_prompt=None,
|
|
history_messages=[],
|
|
**kwargs: openai_complete_if_cache(
|
|
"gpt-4o-mini",
|
|
prompt,
|
|
system_prompt=system_prompt,
|
|
history_messages=history_messages,
|
|
api_key=api_key,
|
|
base_url=base_url,
|
|
**kwargs,
|
|
)
|
|
)
|
|
|
|
|
|
def get_vision_model_func(api_key: str, base_url: str = None):
|
|
return (
|
|
lambda prompt,
|
|
system_prompt=None,
|
|
history_messages=[],
|
|
image_data=None,
|
|
**kwargs: openai_complete_if_cache(
|
|
"gpt-4o",
|
|
"",
|
|
system_prompt=None,
|
|
history_messages=[],
|
|
messages=[
|
|
{"role": "system", "content": system_prompt} if system_prompt else None,
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "text", "text": prompt},
|
|
{
|
|
"type": "image_url",
|
|
"image_url": {
|
|
"url": f"data:image/jpeg;base64,{image_data}"
|
|
},
|
|
},
|
|
],
|
|
}
|
|
if image_data
|
|
else {"role": "user", "content": prompt},
|
|
],
|
|
api_key=api_key,
|
|
base_url=base_url,
|
|
**kwargs,
|
|
)
|
|
if image_data
|
|
else openai_complete_if_cache(
|
|
"gpt-4o-mini",
|
|
prompt,
|
|
system_prompt=system_prompt,
|
|
history_messages=history_messages,
|
|
api_key=api_key,
|
|
base_url=base_url,
|
|
**kwargs,
|
|
)
|
|
)
|
|
|
|
|
|
def get_audio_llm_func(api_key: str, base_url: str = None):
|
|
"""Get audio LLM function for audio analysis"""
|
|
return (
|
|
lambda prompt,
|
|
system_prompt=None,
|
|
history_messages=[],
|
|
**kwargs: openai_complete_if_cache(
|
|
"gpt-4o-audio", # Use audio-capable model if available
|
|
prompt,
|
|
system_prompt=system_prompt,
|
|
history_messages=history_messages,
|
|
api_key=api_key,
|
|
base_url=base_url,
|
|
**kwargs,
|
|
)
|
|
)
|
|
|
|
|
|
async def process_image_example(lightrag: LightRAG, vision_model_func):
|
|
"""Example of processing an image"""
|
|
# Create image processor
|
|
image_processor = ImageModalProcessor(
|
|
lightrag=lightrag, modal_caption_func=vision_model_func
|
|
)
|
|
|
|
# Prepare image content
|
|
image_content = {
|
|
"img_path": "image.jpg",
|
|
"image_caption": ["Example image caption"],
|
|
"image_footnote": ["Example image footnote"],
|
|
}
|
|
|
|
# Process image
|
|
(description, entity_info, _) = await image_processor.process_multimodal_content(
|
|
modal_content=image_content,
|
|
content_type="image",
|
|
file_path="image_example.jpg",
|
|
entity_name="Example Image",
|
|
)
|
|
|
|
print("Image Processing Results:")
|
|
print(f"Description: {description}")
|
|
print(f"Entity Info: {entity_info}")
|
|
|
|
|
|
async def process_table_example(lightrag: LightRAG, llm_model_func):
|
|
"""Example of processing a table"""
|
|
# Create table processor
|
|
table_processor = TableModalProcessor(
|
|
lightrag=lightrag, modal_caption_func=llm_model_func
|
|
)
|
|
|
|
# Prepare table content
|
|
table_content = {
|
|
"table_body": """
|
|
| Name | Age | Occupation |
|
|
|------|-----|------------|
|
|
| John | 25 | Engineer |
|
|
| Mary | 30 | Designer |
|
|
""",
|
|
"table_caption": ["Employee Information Table"],
|
|
"table_footnote": ["Data updated as of 2024"],
|
|
}
|
|
|
|
# Process table
|
|
(description, entity_info, _) = await table_processor.process_multimodal_content(
|
|
modal_content=table_content,
|
|
content_type="table",
|
|
file_path="table_example.md",
|
|
entity_name="Employee Table",
|
|
)
|
|
|
|
print("\nTable Processing Results:")
|
|
print(f"Description: {description}")
|
|
print(f"Entity Info: {entity_info}")
|
|
|
|
|
|
async def process_equation_example(lightrag: LightRAG, llm_model_func):
|
|
"""Example of processing a mathematical equation"""
|
|
# Create equation processor
|
|
equation_processor = EquationModalProcessor(
|
|
lightrag=lightrag, modal_caption_func=llm_model_func
|
|
)
|
|
|
|
# Prepare equation content
|
|
equation_content = {"text": "E = mc^2", "text_format": "LaTeX"}
|
|
|
|
# Process equation
|
|
(description, entity_info, _) = await equation_processor.process_multimodal_content(
|
|
modal_content=equation_content,
|
|
content_type="equation",
|
|
file_path="equation_example.txt",
|
|
entity_name="Mass-Energy Equivalence",
|
|
)
|
|
|
|
print("\nEquation Processing Results:")
|
|
print(f"Description: {description}")
|
|
print(f"Entity Info: {entity_info}")
|
|
|
|
|
|
async def process_audio_example(lightrag: LightRAG, audio_llm_func):
|
|
"""Example of processing audio content"""
|
|
# Create audio processor
|
|
audio_processor = AudioModalProcessor(
|
|
lightrag=lightrag, modal_caption_func=audio_llm_func
|
|
)
|
|
|
|
# Prepare audio content
|
|
audio_content = {
|
|
"audio_path": "./example_audio.mp3", # Example path
|
|
"transcript": "Welcome to this technical tutorial. Today we will explore artificial intelligence and machine learning concepts.",
|
|
"duration": "00:03:45",
|
|
"format": "MP3",
|
|
"description": "AI tutorial introduction audio"
|
|
}
|
|
|
|
# Create example audio file (for demonstration)
|
|
import pathlib
|
|
pathlib.Path("./example_audio.mp3").touch()
|
|
|
|
try:
|
|
# Process audio
|
|
(description, entity_info, _) = await audio_processor.process_multimodal_content(
|
|
modal_content=audio_content,
|
|
content_type="audio",
|
|
file_path="audio_example.mp3",
|
|
entity_name="AI Tutorial Audio",
|
|
)
|
|
|
|
print("\nAudio Processing Results:")
|
|
print(f"Description: {description}")
|
|
print(f"Entity Info: {entity_info}")
|
|
|
|
finally:
|
|
# Clean up example file
|
|
if pathlib.Path("./example_audio.mp3").exists():
|
|
pathlib.Path("./example_audio.mp3").unlink()
|
|
|
|
|
|
async def initialize_rag(api_key: str, base_url: str = None):
|
|
rag = LightRAG(
|
|
working_dir=WORKING_DIR,
|
|
embedding_func=EmbeddingFunc(
|
|
embedding_dim=3072,
|
|
max_token_size=8192,
|
|
func=lambda texts: openai_embed(
|
|
texts,
|
|
model="text-embedding-3-large",
|
|
api_key=api_key,
|
|
base_url=base_url,
|
|
),
|
|
),
|
|
llm_model_func=lambda prompt,
|
|
system_prompt=None,
|
|
history_messages=[],
|
|
**kwargs: openai_complete_if_cache(
|
|
"gpt-4o-mini",
|
|
prompt,
|
|
system_prompt=system_prompt,
|
|
history_messages=history_messages,
|
|
api_key=api_key,
|
|
base_url=base_url,
|
|
**kwargs,
|
|
),
|
|
)
|
|
|
|
await rag.initialize_storages()
|
|
await initialize_pipeline_status()
|
|
|
|
return rag
|
|
|
|
|
|
def main():
|
|
"""Main function to run the example"""
|
|
parser = argparse.ArgumentParser(description="Modal Processors Example")
|
|
parser.add_argument("--api-key", required=True, help="OpenAI API key")
|
|
parser.add_argument("--base-url", help="Optional base URL for API")
|
|
parser.add_argument(
|
|
"--working-dir", "-w", default=WORKING_DIR, help="Working directory path"
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
# Run examples
|
|
asyncio.run(main_async(args.api_key, args.base_url))
|
|
|
|
|
|
async def main_async(api_key: str, base_url: str = None):
|
|
# Initialize LightRAG
|
|
lightrag = await initialize_rag(api_key, base_url)
|
|
|
|
# Get model functions
|
|
llm_model_func = get_llm_model_func(api_key, base_url)
|
|
vision_model_func = get_vision_model_func(api_key, base_url)
|
|
audio_llm_func = get_audio_llm_func(api_key, base_url)
|
|
|
|
# Run examples
|
|
await process_image_example(lightrag, vision_model_func)
|
|
await process_table_example(lightrag, llm_model_func)
|
|
await process_equation_example(lightrag, llm_model_func)
|
|
await process_audio_example(lightrag, audio_llm_func)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|