Spaces:
Sleeping
Sleeping
File size: 2,178 Bytes
c07fa76 5b9ca00 c07fa76 1e72b5a 314c1d9 1e72b5a c07fa76 2bc1685 5b9ca00 c07fa76 1e72b5a 7ca39a5 c07fa76 1e72b5a c07fa76 314c1d9 5b9ca00 c07fa76 1e72b5a c07fa76 1e72b5a c07fa76 5b9ca00 c07fa76 1e72b5a 5b9ca00 c07fa76 a3578e5 1e72b5a | 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 | import os
import tempfile
import requests
from fastapi import HTTPException
# 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
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}")
# splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=120,separators=["\n\n", "\n", ".", " ", ""])
splitter = SentenceTransformersTokenTextSplitter(model_name = MODEL_DIR,tokens_per_chunk=500, chunk_overlap=80)
return splitter.split_documents(docs)
|