Bman21 commited on
Commit
72f5b02
·
verified ·
1 Parent(s): 9ce0803

Create data_loader.py

Browse files
Files changed (1) hide show
  1. data_loader.py +48 -0
data_loader.py ADDED
@@ -0,0 +1,48 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import PyPDF2
3
+ from sentence_transformers import SentenceTransformer
4
+ import faiss
5
+ import numpy as np
6
+
7
+ # Load embedding model
8
+ embedder = SentenceTransformer("all-MiniLM-L6-v2")
9
+
10
+ def load_pdf(file_path):
11
+ """Extract text from a PDF file"""
12
+ text = ""
13
+ with open(file_path, "rb") as f:
14
+ reader = PyPDF2.PdfReader(f)
15
+ for page in reader.pages:
16
+ if page.extract_text():
17
+ text += page.extract_text() + " "
18
+ return text
19
+
20
+ def load_all_pdfs(folder="notes"):
21
+ """Load and merge text from all PDFs in a folder"""
22
+ all_chunks = []
23
+ sources = []
24
+
25
+ for file in os.listdir(folder):
26
+ if file.endswith(".pdf"):
27
+ subject = file.replace(".pdf", "")
28
+ print(f"📖 Loading {subject} ...")
29
+ text = load_pdf(os.path.join(folder, file))
30
+
31
+ # Split into chunks
32
+ chunks = [text[i:i+500] for i in range(0, len(text), 500)]
33
+ all_chunks.extend(chunks)
34
+ sources.extend([subject] * len(chunks)) # Keep track of subject
35
+
36
+ return all_chunks, sources
37
+
38
+ def create_vector_store(chunks):
39
+ """Create embeddings and FAISS index"""
40
+ embeddings = embedder.encode(chunks)
41
+ dim = embeddings.shape[1]
42
+ index = faiss.IndexFlatL2(dim)
43
+ index.add(np.array(embeddings))
44
+ return index
45
+
46
+ # Load all PDFs
47
+ chunks, sources = load_all_pdfs("notes")
48
+ index = create_vector_store(chunks)