Spaces:
Sleeping
Sleeping
| 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)) | |
| splitter = RecursiveCharacterTextSplitter(chunk_size=600, chunk_overlap=80,separators=["\n\n", "\n", ".", " ", ""]) | |
| # splitter = SentenceTransformersTokenTextSplitter(model_name = MODEL_DIR,tokens_per_chunk=500, chunk_overlap=80) | |
| # return splitter.split_documents(docs) | |
| return splitter.split_documents(sentence_docs) | |