uiynyny/llama3RAG
0
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 