File size: 4,585 Bytes
37896f1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 | import os
from dotenv import load_dotenv
from langchain_community.document_loaders import PyMuPDFLoader
from pinecone import Pinecone
from langchain.text_splitter import RecursiveCharacterTextSplitter
from sentence_transformers import SentenceTransformer
import ollama
import numpy as np
from tqdm import tqdm
import hashlib
# PDF_FILE = "DeepSeekR1.pdf"
PDF_FILE = "os.pdf"
# PDF_FILE = "HD_Brochure_2025_26.pdf"
PINECONE_API_KEY = "pcsk_71Vjnk_5MBhhy4u8VajTtHBoJTCULFf2D2FxzssiGgsEqWheg9fAx1nxXFWCynRQtUx384"
INDEX_NAME = "chatdatabase"
# Initialize connections
pc = Pinecone(api_key=PINECONE_API_KEY)
index = pc.Index(INDEX_NAME)
embedder = SentenceTransformer("thenlper/gte-large")
def get_pdf_hash(file_path):
"""Get MD5 hash of PDF file"""
with open(file_path, 'rb') as f:
return hashlib.md5(f.read()).hexdigest()[:12] # Use first 12 chars
def check_embeddings_exist(pdf_hash):
try:
dummy_vector = [0.0] * 1024 # gte-large dimension
results = index.query(
vector=dummy_vector,
filter={"pdf_hash": pdf_hash},
top_k=1
)
return len(results.matches) > 0
except:
return False
def create_chunks(file_path):
"""Load PDF and create chunks"""
print(f"π Loading PDF: {file_path}")
loader = PyMuPDFLoader(file_path)
docs = loader.load()
splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100)
chunks = splitter.split_documents(docs)
print(f"Created {len(chunks)} chunks from {len(docs)} pages")
return chunks
def create_and_store_embeddings(chunks, pdf_hash):
"""Create embeddings and store in Pinecone"""
print("π Creating embeddings...")
vectors = []
for i, chunk in enumerate(tqdm(chunks, desc="Processing chunks")):
# Create embedding
embedding = embedder.encode(chunk.page_content)
# Prepare vector for Pinecone
vectors.append({
"id": f"{pdf_hash}_{i}",
"values": embedding.tolist(),
"metadata": {
"text": chunk.page_content,
"pdf_hash": pdf_hash,
"chunk_id": i,
"page": chunk.metadata.get('page', 0)
}
})
# Upload to Pinecone in batches
print("β¬οΈ Uploading to Pinecone...")
BATCH_SIZE = 100
for i in range(0, len(vectors), BATCH_SIZE):
batch = vectors[i:i+BATCH_SIZE]
index.upsert(vectors=batch)
print(f"Uploaded batch {i//BATCH_SIZE + 1}/{(len(vectors)-1)//BATCH_SIZE + 1}")
print(f"β
Stored {len(vectors)} embeddings")
def query_system(question, top_k=5):
"""Query the system and get answer"""
print(f"\nπ Question: {question}")
# Create embedding for question
question_embedding = embedder.encode(question)
# Search Pinecone
results = index.query(
vector=question_embedding.tolist(),
top_k=top_k,
include_metadata=True
)
if not results.matches:
return "No relevant information found."
context = "\n\n".join([match.metadata['text'] for match in results.matches])
prompt = f"""Based on the following context, answer the question. If the answer is not in the context, say so.
Context:
{context}
Question: {question}
Answer:"""
print("π€ Generating answer...")
try:
response = ollama.generate(model='llama3:latest', prompt=prompt)
answer = response['response']
print(f"π― Answer: {answer}")
return answer
except Exception as e:
print(f"β Error generating answer: {e}")
return "Error generating answer. Check if Ollama is running."
def main():
pdf_hash = get_pdf_hash(PDF_FILE)
print(f"π PDF Hash: {pdf_hash}")
if check_embeddings_exist(pdf_hash):
print("β
Embeddings already exist, skipping creation")
else:
print("π Creating new embeddings...")
chunks = create_chunks(PDF_FILE)
create_and_store_embeddings(chunks, pdf_hash)
print("\nπ System ready! Ask your questions (type 'quit' to exit)")
while True:
question = input("\nβ Your question: ").strip()
if question.lower() in ['quit', 'exit', 'q']:
break
if question:
query_system(question)
if __name__ == "__main__":
main() |