Spaces:
Sleeping
Sleeping
File size: 2,102 Bytes
0879db7 4810f6f 0879db7 4810f6f 0879db7 4810f6f 0879db7 4883530 0879db7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 |
from PIL import Image
import pytesseract
from pdf2image import convert_from_path
import os
from langchain.vectorstores import Chroma
from langchain.embeddings import SentenceTransformerEmbeddings
from langchain.document_loaders import PyMuPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.docstore.document import Document
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import Chroma
from langchain.embeddings import HuggingFaceEmbeddings
CHROMA_DIR = "data/chroma_db"
CHROMA_IMG_DIR = "data/image_db"
def store_pdf(pdf_path):
loader = PyMuPDFLoader(pdf_path)
docs = loader.load()
splitter = RecursiveCharacterTextSplitter(
chunk_size=500, chunk_overlap=100)
chunks = splitter.split_documents(docs)
embeddings = SentenceTransformerEmbeddings(model_name='thenlper/gte-large')
Chroma.from_documents(chunks, embeddings, persist_directory=CHROMA_DIR)
def store_pdf_image(pdf_path):
text = extract_text_from_scanned_pdf(pdf_path)
doc = Document(page_content=text)
splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
chunks = splitter.split_documents([doc])
embeddings = SentenceTransformerEmbeddings(model_name='thenlper/gte-large')
Chroma.from_documents(chunks, embeddings, persist_directory=CHROMA_IMG_DIR)
def store_pdf_image_text(text):
doc = Document(page_content=text)
splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
chunks = splitter.split_documents([doc])
embeddings = SentenceTransformerEmbeddings(model_name='thenlper/gte-large')
Chroma.from_documents(chunks, embeddings, persist_directory=CHROMA_IMG_DIR)
# images = convert_from_path("your_file.pdf", poppler_path="/opt/homebrew/bin")
def extract_text_from_scanned_pdf(pdf_path):
pages = convert_from_path(pdf_path, dpi=300)
all_text = ""
for i, page in enumerate(pages):
text = pytesseract.image_to_string(page, lang="eng")
all_text += f"\n--- Page {i+1} ---\n{text}"
return all_text
|