Spaces:
Sleeping
Sleeping
| import os | |
| from typing import List | |
| from llama_index.core import VectorStoreIndex, Document, StorageContext, load_index_from_storage | |
| from llama_index.embeddings.huggingface import HuggingFaceEmbedding | |
| from langchain.text_splitter import RecursiveCharacterTextSplitter | |
| EMBED_MODEL = "sentence-transformers/all-MiniLM-L6-v2" | |
| DOCS_FOLDER = "data/docs" | |
| SAMPLE_FILE = "data/sample.txt" | |
| STORAGE_DIR = "data/storage" | |
| def load_docs_from_folder(folder_path=DOCS_FOLDER) -> List[Document]: | |
| documents = [] | |
| if not os.path.exists(folder_path): | |
| return documents | |
| for fname in os.listdir(folder_path): | |
| path = os.path.join(folder_path, fname) | |
| if fname.endswith(".txt"): | |
| with open(path, "r", encoding="utf-8") as f: | |
| content = f.read().strip() | |
| if content: | |
| documents.append(Document(text=content)) | |
| elif fname.endswith(".pdf"): | |
| try: | |
| from llama_index.readers.pdf import PDFReader | |
| pdf_docs = PDFReader().load_data(path) | |
| documents.extend([Document(text=doc.text) for doc in pdf_docs]) | |
| except ImportError: | |
| pass | |
| if not documents and os.path.exists(SAMPLE_FILE): | |
| with open(SAMPLE_FILE, "r", encoding="utf-8") as f: | |
| sample_text = f.read().strip() | |
| documents = [Document(text=sample_text)] | |
| if not documents: | |
| documents = [Document(text=" ")] | |
| return documents | |
| def load_vectorstore_from_text(text: str = None, persist: bool = True) -> VectorStoreIndex: | |
| embed_model = HuggingFaceEmbedding(model_name=EMBED_MODEL) | |
| if persist and text is None and os.path.exists(STORAGE_DIR): | |
| storage_context = StorageContext.from_defaults(persist_dir=STORAGE_DIR) | |
| try: | |
| return load_index_from_storage(storage_context) | |
| except Exception as e: | |
| print(f"Failed to load index from storage, rebuilding: {e}") | |
| documents = [Document(text=text.strip())] if text else load_docs_from_folder() | |
| splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=100) | |
| chunk_docs = [] | |
| for doc in documents: | |
| for chunk in splitter.split_text(doc.text): | |
| chunk_docs.append(Document(text=chunk)) | |
| index = VectorStoreIndex.from_documents(chunk_docs, embed_model=embed_model) | |
| if persist: | |
| os.makedirs(STORAGE_DIR, exist_ok=True) | |
| index.storage_context.persist(persist_dir=STORAGE_DIR) | |
| return index | |
| def rebuild_vectorstore(folder_path=DOCS_FOLDER) -> VectorStoreIndex: | |
| documents = load_docs_from_folder(folder_path) | |
| if not documents: | |
| raise ValueError("No documents found in folder!") | |
| return load_vectorstore_from_text(persist=True) | |
| def get_retriever(text: str = None) -> VectorStoreIndex: | |
| try: | |
| return load_vectorstore_from_text(text=text, persist=True) | |
| except Exception as e: | |
| print(f"Vectorstore load failed, rebuilding locally: {e}") | |
| return rebuild_vectorstore() | |