Spaces:
Sleeping
Sleeping
File size: 2,898 Bytes
c07fa76 5b9ca00 c07fa76 70ec4b9 1e72b5a 314c1d9 1e72b5a 12c2b8c c07fa76 70ec4b9 b7c7987 4ddd2aa b7c7987 2bc1685 5b9ca00 c07fa76 1e72b5a 7ca39a5 c07fa76 1e72b5a c07fa76 314c1d9 5b9ca00 c07fa76 1e72b5a c07fa76 1e72b5a c07fa76 5b9ca00 c07fa76 1e72b5a 5b9ca00 c07fa76 2ca1fd2 c07fa76 2ca1fd2 70ec4b9 2ca1fd2 | 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 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | import os
import tempfile
import requests
from fastapi import HTTPException
import nltk
# from langchain_community.document_loaders import PyPDFLoader, Docx2txtLoader
from langchain_community.document_loaders import PyMuPDFLoader, Docx2txtLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_text_splitters.sentence_transformers import SentenceTransformersTokenTextSplitter # give better results but slow can use later for project
from langchain.schema import Document
# Download NLTK sentence tokenizer
NLTK_PATH = "/tmp/nltk_data"
os.makedirs(NLTK_PATH, exist_ok=True)
nltk.data.path.append(NLTK_PATH)
nltk.download("punkt", download_dir=NLTK_PATH, quiet=True)
nltk.download("punkt_tab", download_dir=NLTK_PATH, quiet=True)
MODEL_DIR = os.path.join("/tmp", "e5-large-v2")
def load_and_chunk(url: str) -> list[Document]:
print(url)
resp = requests.get(url)
if resp.status_code != 200:
raise HTTPException(400, "Could not download document")
content_type = resp.headers.get("Content-Type", "").lower()
url_lower = url.lower()
if "application/pdf" in content_type or ".pdf" in url_lower:
with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp:
tmp.write(resp.content)
tmp_path = tmp.name
try:
loader = PyMuPDFLoader(tmp_path)
docs = loader.load_and_split()
finally:
os.remove(tmp_path)
elif (
"application/vnd.openxmlformats-officedocument.wordprocessingml.document" in content_type
or ".docx" in url_lower
):
with tempfile.NamedTemporaryFile(delete=False, suffix=".docx") as tmp:
tmp.write(resp.content)
tmp_path = tmp.name
try:
loader = Docx2txtLoader(tmp_path)
docs = loader.load_and_split()
finally:
os.remove(tmp_path)
elif "text/plain" in content_type or ".txt" in url_lower:
text = resp.content.decode("utf-8", errors="ignore")
docs = [Document(page_content=text)]
else:
raise HTTPException(400, f"Unsupported document type: {content_type}")
# # --- Step 1: Sentence split ---
# sentence_docs = []
# for doc in docs:
# sentences = nltk.sent_tokenize(doc.page_content)
# for sent in sentences:
# if sent.strip():
# sentence_docs.append(Document(page_content=sent, metadata=doc.metadata))
full_text = "\n".join([doc.page_content for doc in docs])
splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=150)
# splitter = SentenceTransformersTokenTextSplitter(model_name = MODEL_DIR,tokens_per_chunk=500, chunk_overlap=80)
# return splitter.split_documents(docs)
# return splitter.split_documents([Document(page_content=full_text)])
return splitter.create_documents([full_text])
|