CoolFace
Apppublic

uiynyny/llama3RAG

sourceHugging Faceupdated 3d agoView on Hugging Face
0likes
EmbeddingManager.py64 linesDownload Raw Back to root
1import os2from langchain_community.document_loaders import UnstructuredPDFLoader3from langchain_pinecone import PineconeVectorStore4from langchain_huggingface import HuggingFaceEmbeddings5from langchain_text_splitters import RecursiveCharacterTextSplitter6from pinecone import Pinecone7 8 9class EmbeddingManager:10    def __init__(11        self,12        model_name: str = "BAAI/bge-small-en",13        device: str = 'cpu',14        encode_kwargs: dict = None,15        host: str = None,16        api_key: str = None,17    ) -> None:18        self.model_name = model_name19        self.device = device20        self.host = host or os.getenv("PINECONE_HOST")21        self.api_key = api_key or os.getenv("PINECONE_API_KEY")22        if not self.host or not self.api_key:23            raise RuntimeError(24                "Provide PINECONE_HOST and PINECONE_API_KEY."25            )26        if encode_kwargs is None:27            encode_kwargs = {"normalize_embeddings": True}28        self.encode_kwargs = encode_kwargs29        self.embeddings = HuggingFaceEmbeddings(30            model_name=self.model_name,31            model_kwargs={"device": self.device},32            encode_kwargs=self.encode_kwargs,33        )34 35    def create_embeddings(self, pdf_path):36        if not os.path.exists(pdf_path):37            raise FileNotFoundError(f"PDF file not found: {pdf_path}")38        loader = UnstructuredPDFLoader(pdf_path)39        docs = loader.load()40        if not docs:41            raise ValueError(f"No documents found in PDF: {pdf_path}")42        text_splitter = RecursiveCharacterTextSplitter(43            chunk_size=1000, chunk_overlap=250)44        splits = text_splitter.split_documents(docs)45        if not splits:46            raise ValueError("No text splits created from the documents.")47        try:48            pinecone = Pinecone(api_key=self.api_key)49            vector_store = PineconeVectorStore(50                index=pinecone.Index(host=self.host),51                embedding=self.embeddings,52            )53            vector_store.add_documents(54                documents=splits,55            )56        except Exception as e:57            raise RuntimeError(f"Error creating vector store: {e}")58        return 059 60 61if __name__ == "__main__":62    manager = EmbeddingManager()63    manager.create_embeddings("0ac948a1fc1a65ea7c7dbcddb232cbaf.pdf")64