Download src/ingestion.py from LightRT/pdf_rag: direct link, hf CLI and curl.
- Browser
- Download file 1.58 kB
-
https://huggingface.co/spaces/LightRT/pdf_rag/resolve/main/src/ingestion.py
- Command line
-
hf download hf://spaces/LightRT/pdf_rag/src/ingestion.py
-
curl -L -o ingestion.py https://huggingface.co/spaces/LightRT/pdf_rag/resolve/main/src/ingestion.py
1.58 kB
| from docling.document_converter import DocumentConverter,PdfFormatOption | |
| from docling.chunking import HybridChunker | |
| from transformers import AutoTokenizer | |
| from langchain_core.documents import Document | |
| from docling.datamodel.base_models import DocumentStream | |
| from io import BytesIO | |
| from docling.datamodel.base_models import InputFormat | |
| from docling.datamodel.pipeline_options import PdfPipelineOptions | |
| pipeline_options = PdfPipelineOptions() | |
| pipeline_options.do_ocr = False | |
| converter = DocumentConverter(format_options={InputFormat.PDF: PdfFormatOption(pipeline_options=pipeline_options)}) | |
| tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2") | |
| chunker = HybridChunker(tokenizer=tokenizer,max_tokens=800,merge_peers=True) | |
| def ingestion_and_chunking(file_bytes: bytes, filename: str) : | |
| source = DocumentStream(name=filename, stream=BytesIO(file_bytes)) | |
| result = converter.convert(source) | |
| chunks = list(chunker.chunk(result.document)) | |
| for chunk in chunks : | |
| chunk.text = chunker.contextualize(chunk) | |
| docs = [] | |
| for chunk in chunks: | |
| pages = sorted({ | |
| prov.page_no | |
| for item in chunk.meta.doc_items | |
| for prov in item.prov}) | |
| docs.append( | |
| Document( | |
| page_content=chunk.text, | |
| metadata={ | |
| "source": chunk.meta.origin.filename, | |
| "pages": pages, | |
| "section": chunk.meta.headings[0] if chunk.meta.headings else None, | |
| } | |
| ) | |
| ) | |
| return docs |