feat(llm): Add Ollama provider and PyMuPDF image extraction
This commit introduces support for Ollama as an alternative Large Language Model (LLM) provider and enhances PDF image extraction capabilities.
- **Ollama Integration:**
- Implemented `set_ollama_config` to configure Ollama's base URL from `config.ini`.
- Modified `llm.py` to dynamically select and configure the LLM (Gemini or Ollama) based on the `PROVIDER` setting.
- Updated `get_model_name` to return provider-specific default model names.
- `pdf_convertor.py` now conditionally initializes `ChatGoogleGenerativeAI` or `ChatOllama` based on the configured provider.
- **PyMuPDF Image Extraction:**
- Added a new `extract_images_from_pdf` function using PyMuPDF (`fitz`) for direct image extraction from PDF files.
- Introduced `get_extract_images_from_pdf_flag` to control this feature via `config.ini`.
- `convert_pdf_to_markdown` and `refine_content` functions were updated to utilize this new image extraction method when enabled.
- **Refinement Flow:**
- Adjusted the order of `save_md_images` in `main.py` and added an option to save the refined markdown with a specific filename (`index_refined.md`).
- **Dependencies:**
- Updated `pyproject.lock` to include new dependencies for Ollama integration (`langchain-ollama`) and PyMuPDF (`PyMuPDF`), along with platform-specific markers for NVIDIA dependencies.
This commit is contained in:
23
refine.py
Normal file
23
refine.py
Normal file
@@ -0,0 +1,23 @@
|
||||
import argparse # New import
|
||||
from pdf_convertor import load_md_file, save_md_images, refine_content
|
||||
from pathlib import Path
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Refine Markdown content from PDF.")
|
||||
parser.add_argument(
|
||||
"--md-path", type=str, required=True, help="Path to the input Markdown file."
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
md_path = args.md_path
|
||||
pdf_path = Path("input").joinpath(Path(args.md_path).parent.name + ".pdf")
|
||||
|
||||
output = Path(md_path).parent
|
||||
output.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
md, images = load_md_file(md_path)
|
||||
with open(pdf_path, "rb") as pdf_file:
|
||||
pdf = pdf_file.read()
|
||||
md = refine_content(md, images, pdf)
|
||||
|
||||
save_md_images(output, md, images, md_name="index_refined.md")
|
||||
Reference in New Issue
Block a user