Spaces:
Sleeping
Sleeping
| import os | |
| import hashlib | |
| from typing import List | |
| import pypdf | |
| from docx import Document | |
| from langchain_text_splitters import RecursiveCharacterTextSplitter | |
| from langchain_community.vectorstores import Chroma | |
| from langchain_huggingface import HuggingFaceEmbeddings | |
| from .config import settings | |
| class DocumentProcessor: | |
| def __init__(self): | |
| self.embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2") | |
| self.vector_store = Chroma( | |
| persist_directory=settings.CHROMA_DB_PATH, | |
| embedding_function=self.embeddings | |
| ) | |
| self.text_splitter = RecursiveCharacterTextSplitter( | |
| chunk_size=settings.CHUNK_SIZE, | |
| chunk_overlap=settings.CHUNK_OVERLAP | |
| ) | |
| def get_file_checksum(self, file_path: str) -> str: | |
| with open(file_path, "rb") as f: | |
| return hashlib.md5(f.read()).hexdigest() | |
| def extract_text(self, file_path: str) -> str: | |
| ext = os.path.splitext(file_path)[1].lower() | |
| text = "" | |
| if ext == ".pdf": | |
| with open(file_path, "rb") as f: | |
| reader = pypdf.PdfReader(f) | |
| for page in reader.pages: | |
| text += page.extract_text() + "\n" | |
| elif ext == ".docx": | |
| doc = Document(file_path) | |
| for para in doc.paragraphs: | |
| text += para.text + "\n" | |
| elif ext == ".txt": | |
| with open(file_path, "r", encoding="utf-8") as f: | |
| text = f.read() | |
| return text | |
| def ingest_document(self, file_path: str, metadata: dict = None): | |
| text = self.extract_text(file_path) | |
| if not text.strip(): | |
| return None | |
| chunks = self.text_splitter.split_text(text) | |
| metadatas = [metadata or {} for _ in chunks] | |
| for meta in metadatas: | |
| meta["source"] = os.path.basename(file_path) | |
| meta["checksum"] = self.get_file_checksum(file_path) | |
| self.vector_store.add_texts(texts=chunks, metadatas=metadatas) | |
| return len(chunks) | |
| def search(self, query: str, k: int = 5): | |
| return self.vector_store.similarity_search(query, k=k) | |
| processor = DocumentProcessor() | |